mirror of
https://github.com/pytorch/pytorch.git
synced 2025-10-20 12:54:11 +08:00
Replace pynvml with nvidia-ml-py in win-test.sh (#164681)
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:
committed by
PyTorch MergeBot
parent
6861fa43e5
commit
48b54b45d6
@ -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
|
||||
|
@ -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
|
||||
|
@ -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
|
||||
|
||||
|
Reference in New Issue
Block a user