Propagate use_orbax_format to _get_optimizer_state_dict.

PiperOrigin-RevId: 611509017
google-research · Feb 29, 2024 · dcbe3e7 · dcbe3e7
1 parent 65ce7e9
commit dcbe3e7
Showing 1 changed file with 8 additions and 2 deletions.
diff --git a/t5x/checkpoints.py b/t5x/checkpoints.py
@@ -1119,7 +1119,9 @@ def restore(
       ckpt_contents = _maybe_update_ts_from_gcs_to_file(ckpt_contents)
 
     ckpt_state_dict = self._get_optimizer_state_dict(
-        ckpt_contents, state_transformation_fns
+        ckpt_contents,
+        state_transformation_fns,
+        use_orbax_format=ckpt_type is checkpoint_utils.CheckpointTypes.ORBAX,
     )
 
     # The state dict may contain TensorStore specs that need to be read.
@@ -1363,9 +1365,13 @@ def _get_optimizer_state_dict(
       self,
       ckpt_contents: PyTree,
       state_transformation_fns: Sequence[RestoreStateTransformationFn],
+      use_orbax_format: bool = False,
   ):
     return _get_optimizer_state_dict(
-        ckpt_contents, self._train_state.state_dict(), state_transformation_fns
+        ckpt_contents,
+        self._train_state.state_dict(),
+        state_transformation_fns,
+        use_orbax_format,
     )