Provide an explanation for import error

Signed-off-by: Mikołaj Błaż <mblaz@nvidia.com>
NVIDIA · Jun 19, 2024 · dbe8423 · dbe8423
1 parent 9549a96
commit dbe8423
Showing 1 changed file with 5 additions and 2 deletions.
diff --git a/nemo/utils/callbacks/dist_ckpt_io.py b/nemo/utils/callbacks/dist_ckpt_io.py
@@ -48,10 +48,13 @@
 
     HAVE_MEGATRON_CORE = True
 
-except (ImportError, ModuleNotFoundError):
+except (ImportError, ModuleNotFoundError) as e:
 
     HAVE_MEGATRON_CORE = False
-    IMPORT_ERROR = "megatron-core was not found. Please see the NeMo README for installation instructions: https://github.com/NVIDIA/NeMo#megatron-gpt."
+    IMPORT_ERROR = (
+        "megatron-core was not found. Please see the NeMo README for installation instructions: https://github.com/NVIDIA/NeMo#megatron-gpt."
+        f" Exact error: {e}"
+    )
 
 
 @contextmanager