Merge pull request #4006 from Uminosachi/scheduler-kwargs

Set scheduler_specific_kwargs to get_scheduler
2024-06-03 19:27:53 +08:00 · 2024-06-03 19:27:53 +08:00 · 0f01500b68
parent 88681d3357 14e97dc119
commit 0f01500b68
1 changed files with 1 additions and 0 deletions
--- 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"):