Update on "Cosmetic changes to train.py"

- Move related code closer - Adjust comments [ghstack-poisoned]
pytorch · Jun 14, 2024 · 36ae0fe · 36ae0fe
2 parents 3002804 + 3aad633
commit 36ae0fe
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/train.py b/train.py
@@ -362,7 +362,7 @@ def loss_fn(pred, labels):
             # optimizer step
             checkpoint.wait_for_staging()
             optimizer.step()
-            scheduler.step()
+            lr_scheduler.step()
 
             losses_since_last_log.append(loss)