add optimizer step to prevent warning in tests

axolotl-ai-cloud · Apr 9, 2024 · 8ee20f3 · 8ee20f3
1 parent 2fa65b9
commit 8ee20f3
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/tests/test_schedulers.py b/tests/test_schedulers.py
@@ -37,11 +37,13 @@ def test_schedulers(self):
         constant_step = int(self.train_steps * self.constant_lr_ratio)
         remaining_step = self.train_steps - constant_step
         for _ in range(constant_step):
+            self.optimizer.step()
             self.lr_scheduler.step()
         self.assertEqual(
             self.lr_scheduler.get_last_lr()[0], self._lr * self.min_lr_ratio
         )
         for _ in range(remaining_step):
+            self.optimizer.step()
             self.lr_scheduler.step()
         self.assertEqual(
             self.lr_scheduler.get_last_lr()[0], self._lr * self.min_lr_ratio