Code update

PiperOrigin-RevId: 623619200
google-research · Apr 10, 2024 · c1bfc9d · c1bfc9d
1 parent 2f5ee6b
commit c1bfc9d
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 3 deletions.
diff --git a/swirl_dynamics/templates/callbacks.py b/swirl_dynamics/templates/callbacks.py
@@ -300,8 +300,9 @@ class ParameterOverview(Callback):
   log_to_tb: bool = True
 
   def on_train_begin(self, trainer: trainers.BaseTrainer) -> None:
-    if isinstance(trainer.train_state, train_states.BasicTrainState):
-      params = trainer.train_state.params
+    train_state = trainer.unreplicated_train_state
+    if isinstance(train_state, train_states.BasicTrainState):
+      params = train_state.params
       if self.log_to_info:
         logging.info("Logging parameter overview.")
         parameter_overview.log_parameter_overview(params)

diff --git a/swirl_dynamics/templates/callbacks_test.py b/swirl_dynamics/templates/callbacks_test.py
@@ -180,7 +180,7 @@ def test_logs_parameter_overview(self):
     callback = callbacks.ParameterOverview()
     callback.metric_writer = metric_writers.create_default_writer(work_dir)
     trainer = mock.Mock(spec=trainers.BaseTrainer)
-    trainer.train_state = train_states.BasicTrainState(
+    trainer.unreplicated_train_state = train_states.BasicTrainState(
         step=jnp.array(0),
         params={"bias": jnp.ones((10,)), "weights": jnp.ones((10, 10))},
         opt_state={},