Ensure that all steps written are real integers instead of JAX device arrays with type int.
Ensure that all steps written are real integers instead of JAX device arrays with type int.

PiperOrigin-RevId: 426413295

google-research/t5x