Update TfGrainCheckpointHandler.

PiperOrigin-RevId: 686405738
google · Oct 17, 2024 · 0286a94 · 0286a94
1 parent b9d4ca5
commit 0286a94
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/ffn/jax/train.py b/ffn/jax/train.py
@@ -432,7 +432,7 @@ def train_and_evaluate(
         train_state_path, args=ocp.args.StandardRestore(state)
     )
     checkpointed_state['train_iter'] = iter_handler.restore(
-        train_iter_path, args
+        train_iter_path, args=args
     )
     logging.info('Initializing training from %r', config.init_from_cpoint)
   elif latest_step is not None: