Revert "Fix import guard checks (NVIDIA#7124)" (NVIDIA#7125)

This reverts commit ae7624d.
rohitrango · Jul 28, 2023 · 1a70ed4 · 1a70ed4
1 parent ae7624d
commit 1a70ed4
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 15 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, AttributeError):
+    except (ImportError, ModuleNotFoundError):
         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,23 +44,16 @@
 except (ImportError, ModuleNotFoundError):
     HAVE_TRITON = False
 
-try:
-    import pynvml
-
-    HAVE_PYNVML = True
-except (ImportError, ModuleNotFoundError):
-    HAVE_PYNVML = False
+import pynvml
 
 
 def HAVE_AMPERE_GPU():
-    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
+    pynvml.nvmlInit()
+    handle = pynvml.nvmlDeviceGetHandleByIndex(0)
+    device_arch = pynvml.nvmlDeviceGetArchitecture(handle)
+    pynvml.nvmlShutdown()
+    return device_arch == pynvml.NVML_DEVICE_ARCH_AMPERE
+
 
 @pytest.mark.run_only_on('GPU')
 @pytest.mark.skipif(not HAVE_APEX, reason="apex is not installed")