grr

stanford-crfm · Jan 22, 2025 · 7cd4791 · 7cd4791
1 parent bc32f43
commit 7cd4791
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/src/levanter/trainer_state.py b/src/levanter/trainer_state.py
@@ -53,7 +53,7 @@ class TrainerState(eqx.Module, Generic[M]):
     optimizer: GradientTransformation = eqx.field(static=True)
     opt_state: OptState
     training_key: PRNGKeyArray
-    model_averaging: ModelAveraging[M] | None
+    model_averaging: ModelAveraging[M] | None = None
 
     is_trainable: FilterTree = eqx.field(static=True)
     mp: jmp.Policy = eqx.field(static=True)