[CI][Build] change pynvml to nvidia-ml-py (vllm-project#4302)

esmeetu · Apr 24, 2024 · e4bf860 · e4bf860
1 parent 91f50a6
commit e4bf860
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/requirements-cuda.txt b/requirements-cuda.txt
@@ -3,7 +3,7 @@
 
 # Dependencies for NVIDIA GPUs
 ray >= 2.9
-pynvml == 11.5.0
+nvidia-ml-py # for pynvml package
 vllm-nccl-cu12>=2.18,<2.19  # for downloading nccl library
 torch == 2.2.1
 xformers == 0.0.25  # Requires PyTorch 2.2.1