make sure to create scheduler first

axolotl-ai-cloud · Jun 4, 2024 · 8e720d8 · 8e720d8
1 parent 263a351
commit 8e720d8
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
@@ -390,6 +390,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,