make sure to create scheduler first

axolotl-ai-cloud · Jul 31, 2024 · 78d4a1e · 78d4a1e
1 parent 27a78f9
commit 78d4a1e
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/src/axolotl/core/trainer_builder.py b/src/axolotl/core/trainer_builder.py
@@ -372,6 +372,7 @@ def create_scheduler(
                 anneal_steps = (
                     self.args.jagged_restarts_anneal_steps or 1
                 )
+                super().create_scheduler(num_training_steps, optimizer)
                 self.lr_scheduler = JaggedLRRestartScheduler(  # pylint: disable=attribute-defined-outside-init
                     optimizer,
                     self.lr_scheduler,