frozenleaves/pytorch
1
0
Fork 0
mirror of https://github.com/pytorch/pytorch.git synced 2025-10-21 05:34:18 +08:00
Code Issues Packages Projects Releases Wiki Activity

Replace pynvml with nvidia-ml-py in win-test.sh (#164681)

Browse Source 
pynvml was deprecated.

Pull Request resolved: https://github.com/pytorch/pytorch/pull/164681
Approved by: https://github.com/Aidyn-A, https://github.com/eqy
This commit is contained in:
Yuanyuan Chen
2025-10-06 21:57:26 +00:00
committed by PyTorch MergeBot
parent 6861fa43e5
commit 48b54b45d6
3 changed files with 3 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
.ci/pytorch/win-test.sh
View File

@ -38,7 +38,7 @@ if [[ "$BUILD_ENVIRONMENT" == *cuda* ]]; then
fi
# TODO: Move this to .ci/docker/requirements-ci.txt
python -m pip install "psutil==5.9.1" "pynvml==11.4.1" "pytest-shard==0.1.2"
python -m pip install "psutil==5.9.1" nvidia-ml-py "pytest-shard==0.1.2"
run_tests() {
# Run nvidia-smi if available

2
torch/cuda/__init__.py
View File

@ -1210,7 +1210,7 @@ def get_sync_debug_mode() -> int:
def _get_pynvml_handler(device: "Device" = None):
if not _HAS_PYNVML:
raise ModuleNotFoundError(
"pynvml does not seem to be installed or it can't be imported."
"nvidia-ml-py does not seem to be installed or it can't be imported."
# pyrefly: ignore # invalid-inheritance
) from _PYNVML_ERR
# pyrefly: ignore # import-error

2
torch/cuda/memory.py
View File

@ -769,7 +769,7 @@ def list_gpu_processes(device: "Device" = None) -> str:
try:
import pynvml # type: ignore[import]
except ModuleNotFoundError:
return "pynvml module not found, please install pynvml"
return "pynvml module not found, please install nvidia-ml-py"
# pyrefly: ignore # import-error
from pynvml import NVMLError_DriverNotLoaded