Merge pull request #4006 from Uminosachi/scheduler-kwargs

Set scheduler_specific_kwargs to get_scheduler
hiyouga · Jun 3, 2024 · 0f01500 · 0f01500
2 parents 88681d3 + 14e97dc
commit 0f01500
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/src/llamafactory/train/utils.py b/src/llamafactory/train/utils.py
@@ -379,6 +379,7 @@ def create_custom_scheduler(
                 optimizer=optimizer_dict[param],
                 num_warmup_steps=training_args.get_warmup_steps(num_training_steps),
                 num_training_steps=num_training_steps,
+                scheduler_specific_kwargs=training_args.lr_scheduler_kwargs,
             )
 
         def scheduler_hook(param: "torch.nn.Parameter"):