Fix issue with restoring orbax checkpoints through transformation_fn

PiperOrigin-RevId: 576571694
google-research · Oct 25, 2023 · 8a294fb · 8a294fb
1 parent 0a56776
commit 8a294fb
Showing 1 changed file with 5 additions and 1 deletion.
diff --git a/t5x/checkpoints.py b/t5x/checkpoints.py
@@ -2361,7 +2361,11 @@ def restore(
             self.directory,
             restore_dtype=self._restore_dtype,
         )
-        return legacy_checkpointer.restore(path=path)
+        return legacy_checkpointer.restore(
+            path=path,
+            fallback_state=fallback_state,
+            state_transformation_fns=state_transformation_fns,
+        )
 
     state_dict = self._train_state.state_dict()
     # Returns a state dict rather than a train state.