Add lr_scheduler check for inference.

Signed-off-by: Dashiell Stander <[email protected]>
EleutherAI · Oct 17, 2023 · a97bd1f · a97bd1f
1 parent f6ac04d
commit a97bd1f
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/megatron/training.py b/megatron/training.py
@@ -686,7 +686,9 @@ def setup_model_and_optimizer(neox_args, use_cache=False, iteration=None):
         neox_args.iteration = 0
 
     # need this for correct lr scheduling resume from ckpt
-    lr_scheduler.optimizer = model.optimizer
+    # but it will not exist if this is being called for inference
+    if lr_scheduler is not None:
+        lr_scheduler.optimizer = model.optimizer
 
     return model, optimizer, lr_scheduler