Fix a bug where a refactoring cl missed passing step when restoring a checkpoint, causing spmd decode once to always use the latest checkpoint and ignore the explicit step specified. (pmap is not affected.)
Fix a bug where a refactoring cl missed passing step when restoring a checkpoint, causing spmd decode once to always use the latest checkpoint and ignore the explicit step specified. (pmap is not affected.)

PiperOrigin-RevId: 459282769

google/paxml,google/paxml