update trainer

stanford-crfm · Sep 11, 2024 · 45f1558 · 45f1558
1 parent 12b55c2
commit 45f1558
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/src/levanter/trainer.py b/src/levanter/trainer.py
@@ -339,8 +339,8 @@ def init_state_and_model(model_init, training_key):
         trainer_state_shape = eqx.filter_eval_shape(init_state_and_model, model_init, training_key)
 
         saveable_train_state = saveable_training_mask(trainer_state_shape, is_trainable)
-        if self.config.reset_optimizer_state:
-            saveable_train_state = dataclasses.replace(saveable_train_state, optimizer=False)
+        # if self.config.reset_optimizer_state:
+        #     saveable_train_state = dataclasses.replace(saveable_train_state, optimizer=False)
 
         state = load_checkpoint_or_initialize(
             init_state_and_model,