[train_unconditional] fix LR scheduler init (#2010)

fix lr scheduler

[train_unconditional] fix LR scheduler init (#2010)
fix lr scheduler
f861cde1 · Suraj Patil · GitHub · b2ea8a84 · f861cde1
Unverified Commit f861cde1 authored Jan 17, 2023 by Suraj Patil Committed by GitHub Jan 17, 2023
Hide whitespace changes
Inline Side-by-side

Showing with 2 additions and 2 deletions

examples/unconditional_image_generation/train_unconditional.py ...les/unconditional_image_generation/train_unconditional.py +2 -2

No files found.
--- a/examples/unconditional_image_generation/train_unconditional.py
+++ b/examples/unconditional_image_generation/train_unconditional.py
@@ -338,8 +338,8 @@ def main(args):
    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(