Consider pynvml available when installed through the nvidia-ml-py distribution

matthewdouglas · matthewdouglas · commit e4f1bf3d0a8b · 2024-07-15T18:59:08.000-04:00
diff --git a/src/accelerate/utils/imports.py b/src/accelerate/utils/imports.py
@@ -86,7 +86,7 @@ def is_import_timer_available():
 
 
 def is_pynvml_available():
-    return _is_package_available("pynvml")
+    return _is_package_available("pynvml") or _is_package_available("pynvml", "nvidia-ml-py")
 
 
 def is_pytest_available():
