Consider pynvml available when installed through the nvidia-ml-py dis…

…tribution (#2936)
huggingface · Jul 16, 2024 · 7a5c231 · 7a5c231
1 parent 4f02bb7
commit 7a5c231
Showing 1 changed file with 1 addition and 1 deletion.
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():