Internal change.

PiperOrigin-RevId: 592911220
kkpan11 · Dec 21, 2023 · 133d0b1 · 133d0b1
1 parent f54468c
commit 133d0b1
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/t5x/checkpoints.py b/t5x/checkpoints.py
@@ -2249,7 +2249,7 @@ class OrbaxCheckpointManagerInterface:
   class _CheckpointManagerImpl(ocp.CheckpointManager):
     """CheckpointManager implementation to deal with metrics update."""
 
-    def _get_old_steps_to_remove(self):
+    def _get_old_steps_to_remove(self) -> List[int]:
       """Update metrics for Orbax management, if available."""
       if self._track_best:
         metric_name_to_monitor = self._options.metric_name_to_monitor  # pytype: disable=attribute-error
@@ -2262,7 +2262,7 @@ def _get_old_steps_to_remove(self):
           if info.step in step_to_metric:
             metrics = {metric_name_to_monitor: step_to_metric[info.step]}
             info.metrics = metrics
-      super()._get_old_steps_to_remove()
+      return super()._get_old_steps_to_remove()
 
   def __init__(
       self,