Trax training loop stops after 1 epoch even after setting n_steps to 10

Question

https://colab.research.google.com/drive/1Vfc5i2NkqHdCHwnOHRFP7nr40y2kRA3_?usp=sharing

Please run the linked notebook. My trax training loop stops after 1 epoch without any error messages. Please advice where am I making a mistake.

Relevant Code bits