how'd i miess this

stanford-crfm · Nov 19, 2024 · c8c2021 · c8c2021
1 parent 516be5e
commit c8c2021
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/src/levanter/optim/config.py b/src/levanter/optim/config.py
@@ -172,7 +172,7 @@ def lr_scheduler(self, num_train_steps):
             if warmup_steps != 0:
                 warmup = optax.linear_schedule(previous_end, self.learning_rate, warmup_steps)
                 schedules.append(warmup)
-                boundaries.append(warmup_steps)
+                boundaries.append(start + warmup_steps)
 
             stable_steps = _convert_ratio_or_steps(self.stable, cycle_steps)
             lr_decay_steps = cycle_steps - stable_steps - warmup_steps