ljkdmald

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