mirror of
https://github.com/pytorch/pytorch.git
synced 2025-11-03 23:45:05 +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
|
fi
|
||||||
|
|
||||||
# TODO: Move this to .ci/docker/requirements-ci.txt
|
# 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_tests() {
|
||||||
# Run nvidia-smi if available
|
# Run nvidia-smi if available
|
||||||
|
|||||||
@ -1210,7 +1210,7 @@ def get_sync_debug_mode() -> int:
|
|||||||
def _get_pynvml_handler(device: "Device" = None):
|
def _get_pynvml_handler(device: "Device" = None):
|
||||||
if not _HAS_PYNVML:
|
if not _HAS_PYNVML:
|
||||||
raise ModuleNotFoundError(
|
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
|
# pyrefly: ignore # invalid-inheritance
|
||||||
) from _PYNVML_ERR
|
) from _PYNVML_ERR
|
||||||
# pyrefly: ignore # import-error
|
# pyrefly: ignore # import-error
|
||||||
|
|||||||
@ -769,7 +769,7 @@ def list_gpu_processes(device: "Device" = None) -> str:
|
|||||||
try:
|
try:
|
||||||
import pynvml # type: ignore[import]
|
import pynvml # type: ignore[import]
|
||||||
except ModuleNotFoundError:
|
except ModuleNotFoundError:
|
||||||
return "pynvml module not found, please install pynvml"
|
return "pynvml module not found, please install nvidia-ml-py"
|
||||||
# pyrefly: ignore # import-error
|
# pyrefly: ignore # import-error
|
||||||
from pynvml import NVMLError_DriverNotLoaded
|
from pynvml import NVMLError_DriverNotLoaded
|
||||||
|
|
||||||
|
|||||||
Reference in New Issue
Block a user