Internal change

recml/core/training/jax_trainer.py

@@ -455,8 +455,10 @@ def checkpoint_manager(self) -> ocp.CheckpointManager:
 
     save_on_steps = set(save_on_steps)
 
+    checkpoint_dir = os.path.join(self._model_dir, core.CHECKPOINT_DIR)
+
     return ocp.CheckpointManager(
-        directory=os.path.join(self._model_dir, core.CHECKPOINT_DIR),
+        directory=checkpoint_dir,
         options=ocp.CheckpointManagerOptions(
             should_save_fn=lambda step, _: step in save_on_steps,
             max_to_keep=self._max_checkpoints_to_keep,

recml/core/training/jax_trainer_test.py

@@ -16,9 +16,11 @@
 from collections.abc import Mapping, Sequence
 import dataclasses
 import os
+from unittest import mock
 
 from absl import flags
 from absl.testing import absltest
+from absl.testing import flagsaver
 from absl.testing import parameterized
 import clu.metrics as clu_metrics
 import flax.linen as nn
@@ -260,6 +262,5 @@ class State:
     self.assertIn("optimizer/learning_rate", metrics)
     self.assertEqual(metrics["optimizer/learning_rate"].compute(), 0.1)
 
-
 if __name__ == "__main__":
   absltest.main()

recml/core/utils/keras_utils_test.py

@@ -23,6 +23,7 @@
 
 from absl import flags
 from absl.testing import absltest
+from absl.testing import flagsaver
 from absl.testing import parameterized
 import jax
 import jax.numpy as jnp