Fix import guard checks (NVIDIA#7124)

Signed-off-by: smajumdar <titu1994@gmail.com>
dorotat-nv · Jul 28, 2023 · a46e325 · a46e325
1 parent fcfc0eb
commit a46e325
Show file tree

Hide file tree

Showing 2 changed files with 15 additions and 8 deletions.
diff --git a/nemo/utils/model_utils.py b/nemo/utils/model_utils.py
@@ -576,7 +576,7 @@ def check_lib_version(lib_name: str, checked_version: str, operator) -> Tuple[Op
                 f"Could not check version compatibility."
             )
             return False, msg
-    except (ImportError, ModuleNotFoundError):
+    except (ImportError, ModuleNotFoundError, AttributeError):
         pass
 
     msg = f"Lib {lib_name} has not been installed. Please use pip or conda to install this package."

diff --git a/tests/collections/nlp/test_flash_attention.py b/tests/collections/nlp/test_flash_attention.py
@@ -44,16 +44,23 @@
 except (ImportError, ModuleNotFoundError):
     HAVE_TRITON = False
 
-import pynvml
+try:
+    import pynvml
+
+    HAVE_PYNVML = True
+except (ImportError, ModuleNotFoundError):
+    HAVE_PYNVML = False
 
 
 def HAVE_AMPERE_GPU():
-    pynvml.nvmlInit()
-    handle = pynvml.nvmlDeviceGetHandleByIndex(0)
-    device_arch = pynvml.nvmlDeviceGetArchitecture(handle)
-    pynvml.nvmlShutdown()
-    return device_arch == pynvml.NVML_DEVICE_ARCH_AMPERE
-
+    if HAVE_PYNVML:
+        pynvml.nvmlInit()
+        handle = pynvml.nvmlDeviceGetHandleByIndex(0)
+        device_arch = pynvml.nvmlDeviceGetArchitecture(handle)
+        pynvml.nvmlShutdown()
+        return device_arch == pynvml.NVML_DEVICE_ARCH_AMPERE
+    else:
+        return False
 
 @pytest.mark.run_only_on('GPU')
 @pytest.mark.skipif(not HAVE_APEX, reason="apex is not installed")