Skip to content

Commit

Permalink
[train_unconditional] fix LR scheduler init (huggingface#2010)
Browse files Browse the repository at this point in the history
fix lr scheduler
  • Loading branch information
patil-suraj authored Jan 17, 2023
1 parent b2ea8a8 commit f861cde
Showing 1 changed file with 2 additions and 2 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -338,8 +338,8 @@ def transforms(examples):
lr_scheduler = get_scheduler(
args.lr_scheduler,
optimizer=optimizer,
num_warmup_steps=args.lr_warmup_steps,
num_training_steps=(len(train_dataloader) * args.num_epochs) // args.gradient_accumulation_steps,
num_warmup_steps=args.lr_warmup_steps * args.gradient_accumulation_steps,
num_training_steps=(len(train_dataloader) * args.num_epochs),
)

model, optimizer, train_dataloader, lr_scheduler = accelerator.prepare(
Expand Down

0 comments on commit f861cde

Please sign in to comment.