oops

stanford-crfm · Jan 23, 2025 · 1936d28 · 1936d28
1 parent aa25122
commit 1936d28
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/src/levanter/trainer_state.py b/src/levanter/trainer_state.py
@@ -115,7 +115,7 @@ def init(
         if model_averaging is not None:
             model_averaging = model_averaging.create(trainable_model)
 
-        opt_state = init_optimizer_for_trainables(optimizer, model)
+        opt_state = init_optimizer_for_trainables(optimizer, trainable_model)
 
         return cls(
             0,