Update config.py

stanford-crfm · Nov 21, 2024 · a88caab · a88caab
1 parent f44e5c8
commit a88caab
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/src/levanter/optim/config.py b/src/levanter/optim/config.py
@@ -180,7 +180,7 @@ def lr_scheduler(self, num_train_steps):
             if stable_steps != 0:
                 stable = optax.constant_schedule(self.learning_rate)
                 schedules.append(stable)
-                boundaries.append(start + warmup_steps + stable_steps)
+                boundaries.append(start + stable_steps)
 
             match self.lr_schedule:
                 case "constant":