fix optimizer number of steps #3428

✅ All contributors are covered under a CLA with Google

See https://cla.developers.google.com/ for more info about Google's Contributor License Agreement (CLA).

ℹ️ Googlers: Go here to view more details and manage scans for this pull request.

@entrpn

The following contributors were found for this pull request:

✅ e6597bb Author: @entrpn <jua*****do@gmail.com>
✅ 6e51359 Author: @aireenmei <ai*****ei@gmail.com>
✅ 093ab89 Author: @Google-ML-Automation <googl**********ation@google.com>
✅ c6b84c1 Author: @jesselu-google <je****u@google.com>
✅ cf051eb Author: @shuningjin <shu*****in@google.com>
✅ 970808c Author: @igorts-git <ig***s@google.com>

(Only the first commit for a unique contributor is listed.)