Fix the error that's happening in T5X OSS github actions. Previously, this test was logging the stale train_state which got donated.
Fix the error that's happening in T5X OSS github actions. Previously, this test was logging the stale train_state which got donated.

PiperOrigin-RevId: 486764305

google-research/t5x