no more cb_states

stanford-crfm · Nov 24, 2024 · c8553b4 · c8553b4
1 parent f181656
commit c8553b4
Showing 1 changed file with 1 addition and 2 deletions.
diff --git a/tests/test_checkpoint.py b/tests/test_checkpoint.py
@@ -41,7 +41,6 @@ def _dummy_step_info(step):
             training_key=jax.random.PRNGKey(0),
             is_trainable=True,
             mp=None,
-            cb_states=(),
         ),
         loss=0.0,
         step_duration=0.0,
@@ -168,7 +167,7 @@ def _make_state(step, key):
     optim = optax.adam(1e-4)
     opt_state = optim.init(arrays_only(model))
 
-    return TrainerState(step, model, optim, opt_state, key, is_trainable=True, mp=None, cb_states=())
+    return TrainerState(step, model, optim, opt_state, key, is_trainable=True, mp=None)
 
 
 def test_checkpoint_simple():