Set scheduler_specific_kwargs to get_scheduler

Former-commit-id: 14e97dc1192f6cf94ab99eb3a9b8c64029040384
This commit is contained in:
Uminosachi 2024-05-31 13:45:39 +09:00
parent 72ebcb9a04
commit 0de4e1e9e2

View File

@ -379,6 +379,7 @@ def create_custom_scheduler(
optimizer=optimizer_dict[param], optimizer=optimizer_dict[param],
num_warmup_steps=training_args.get_warmup_steps(num_training_steps), num_warmup_steps=training_args.get_warmup_steps(num_training_steps),
num_training_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"): def scheduler_hook(param: "torch.nn.Parameter"):