fix ckpt load bug (#9621)

* fix ckpt load bug Signed-off-by: dimapihtar <[email protected]> * Apply isort and black reformatting Signed-off-by: dimapihtar <[email protected]> --------- Signed-off-by: dimapihtar <[email protected]> Signed-off-by: dimapihtar <[email protected]> Co-authored-by: dimapihtar <[email protected]> Signed-off-by: Tugrul Konuk <[email protected]>
NVIDIA · Jul 19, 2024 · 82c529f · 82c529f
1 parent 544b8e8
commit 82c529f
Showing 1 changed file with 2 additions and 2 deletions.
diff --git 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