fix: num training steps for lr decay

1 year ago · 9dfd8e1a7c
parent 195f8a7d4e
commit 9dfd8e1a7c
1 changed files with 1 additions and 1 deletions
--- a/train.py
+++ b/train.py
@ -100,7 +100,7 @@ def train(accelerator, config):
            name="cosine",
            optimizer=optimizer,
            num_warmup_steps=config["warmup_steps"] * accelerator.num_processes,
-            num_training_steps=total_num_steps * accelerator.num_processes,
+            num_training_steps=total_num_steps,
        )
    else:
        scheduler = DummyScheduler(