From 0c8654fef6d95053db21e39f96df59cc5d7a253b Mon Sep 17 00:00:00 2001
From: Rohan138 <rohanpotdar138@gmail.com>
Date: Mon, 5 Aug 2024 23:37:11 +0000
Subject: [PATCH] move pynvml install to setup.py

---
 requirements/requirements.txt | 1 -
 setup.py                      | 4 ++++
 2 files changed, 4 insertions(+), 1 deletion(-)

diff --git a/requirements/requirements.txt b/requirements/requirements.txt
index 05f88337f3a9..6840d6dbcc98 100755
--- a/requirements/requirements.txt
+++ b/requirements/requirements.txt
@@ -1,7 +1,6 @@
 hjson
 ninja
 numpy
-nvidia-ml-py
 packaging>=20.0
 psutil
 py-cpuinfo
diff --git a/setup.py b/setup.py
index 2b7555361655..8707209526ad 100755
--- a/setup.py
+++ b/setup.py
@@ -92,6 +92,10 @@ def get_env_if_set(key, default: typing.Any = ""):
     'triton': fetch_requirements('requirements/requirements-triton.txt'),
 }
 
+# Only install pynvml on nvidia gpus.
+if torch_available and get_accelerator().device_name() == 'cuda' and not is_rocm_pytorch:
+    install_requires.append('nvidia-ml-py')
+
 # Add specific cupy version to both onebit extension variants.
 if torch_available and get_accelerator().device_name() == 'cuda':
     cupy = None