From a97bd1fa4ed77929bea730da4de40d6001291b3d Mon Sep 17 00:00:00 2001 From: Dashiell Stander Date: Tue, 17 Oct 2023 02:49:10 -0400 Subject: [PATCH] Add lr_scheduler check for inference. Signed-off-by: Dashiell Stander --- megatron/training.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/megatron/training.py b/megatron/training.py index 0f9f7a0e0..548f81cb0 100644 --- 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