fix import (NVIDIA#9240)

* fix import Signed-off-by: Alexandros Koumparoulis <[email protected]> * Apply isort and black reformatting Signed-off-by: akoumpa <[email protected]> --------- Signed-off-by: Alexandros Koumparoulis <[email protected]> Signed-off-by: akoumpa <[email protected]> Co-authored-by: akoumpa <[email protected]>
rohitrango · May 20, 2024 · 31a23f5 · 31a23f5
1 parent e77d693
commit 31a23f5
Showing 1 changed file with 4 additions and 4 deletions.
diff --git a/examples/nlp/language_modeling/megatron_lm_ckpt_to_nemo.py b/examples/nlp/language_modeling/megatron_lm_ckpt_to_nemo.py
@@ -291,9 +291,9 @@ def load_from_checkpoint(
     **kwargs,
 ):
     """
-        Loads Megatron_LM checkpoints, convert it, with some maintenance of restoration.
-        For documentation, please refer to LightningModule.load_from_checkpoin() documentation.
-        """
+    Loads Megatron_LM checkpoints, convert it, with some maintenance of restoration.
+    For documentation, please refer to LightningModule.load_from_checkpoin() documentation.
+    """
     checkpoint = None
     try:
         cls._set_model_restore_state(is_being_restored=True)
@@ -470,7 +470,7 @@ def convert(local_rank, rank, world_size, args):
             )
         if mcore_output and not args.mcore_input:
             # convert from legacy Megatron-LM to MCore NeMo. Initialize an mcore translation dict
-            from scripts.nlp_language_modeling.convert_nemo_gpt_to_mcore import build_key_mapping
+            from scripts.checkpoint_converters.convert_gpt_nemo_to_mcore import build_key_mapping
 
             mcore_translate = {}
             for k, v in build_key_mapping(model_cfg).items():