From 82c529feed7c5daada17f5996438193fbe4ab2f5 Mon Sep 17 00:00:00 2001 From: Dmytro Pykhtar <37850217+dimapihtar@users.noreply.github.com> Date: Sat, 6 Jul 2024 03:05:03 +0300 Subject: [PATCH] fix ckpt load bug (#9621) * fix ckpt load bug Signed-off-by: dimapihtar * Apply isort and black reformatting Signed-off-by: dimapihtar --------- Signed-off-by: dimapihtar Signed-off-by: dimapihtar Co-authored-by: dimapihtar Signed-off-by: Tugrul Konuk --- nemo/collections/nlp/parts/nlp_overrides.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/nemo/collections/nlp/parts/nlp_overrides.py b/nemo/collections/nlp/parts/nlp_overrides.py index ad220aaa3539..0b89bfda8dbd 100644 --- a/nemo/collections/nlp/parts/nlp_overrides.py +++ b/nemo/collections/nlp/parts/nlp_overrides.py @@ -573,8 +573,8 @@ def _integrate_original_checkpoint_data(self, checkpoint: Dict[str, Any]) -> Dic ]['optimizer']['param_groups'] else: checkpoint['optimizer_states'][0]['param_groups'] = original_checkpoint['optimizer_states'][0][ - 'optimizer' - ]['param_groups'] + 'param_groups' + ] return checkpoint