missed a spot

stanford-crfm · Nov 25, 2024 · c9e2f97 · c9e2f97
1 parent a1c264e
commit c9e2f97
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/tests/test_eval_lm.py b/tests/test_eval_lm.py
@@ -35,7 +35,7 @@ def test_eval_lm():
             Vocab = haliax.Axis("vocab", len(tok))
             model = Gpt2LMHeadModel.init(Vocab, model_config, key=jax.random.PRNGKey(0))
 
-            state = TrainerState(0, model, model, jax.random.PRNGKey(0), True, None, None, ())
+            state = TrainerState(0, model, model, jax.random.PRNGKey(0), True, None, None)
 
             save_checkpoint(state, 0, f"{f}/ckpt")