Make OrbaxCheckpointManagerInterface gin.configurable.

PiperOrigin-RevId: 636264343
google-research · May 22, 2024 · 7bd9794 · 7bd9794
1 parent 3ea9b2c
commit 7bd9794
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/t5x/checkpoints.py b/t5x/checkpoints.py
@@ -2264,6 +2264,7 @@ def _partition_parameter(maybe_arr: Any, param_info: _OrbaxParamInfo):
   return train_state.restore_state(state_dict)
 
 
+@gin.configurable
 class OrbaxCheckpointManagerInterface:
   """Wrapper for ocp.CheckpointManager."""