From 353546614b99872cc4b16183d4db1f7878d48a82 Mon Sep 17 00:00:00 2001 From: RecML authors Date: Fri, 30 Oct 2026 09:33:03 -0700 Subject: [PATCH] Internal change PiperOrigin-RevId: 1010576275 --- recml/core/training/jax_trainer.py | 4 +++- recml/core/training/jax_trainer_test.py | 3 ++- recml/core/utils/keras_utils_test.py | 1 + 3 files changed, 6 insertions(+), 2 deletions(-) diff --git a/recml/core/training/jax_trainer.py b/recml/core/training/jax_trainer.py index 0a73141..7227358 100644 --- a/recml/core/training/jax_trainer.py +++ b/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, diff --git a/recml/core/training/jax_trainer_test.py b/recml/core/training/jax_trainer_test.py index 701e7dd..aac5e07 100644 --- a/recml/core/training/jax_trainer_test.py +++ b/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() diff --git a/recml/core/utils/keras_utils_test.py b/recml/core/utils/keras_utils_test.py index 88ae6bb..6209588 100644 --- a/recml/core/utils/keras_utils_test.py +++ b/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