diff --git a/pkgs/development/python-modules/nvidia-ml-py/default.nix b/pkgs/development/python-modules/nvidia-ml-py/default.nix index 95604dcb20d7..08beac0daf50 100644 --- a/pkgs/development/python-modules/nvidia-ml-py/default.nix +++ b/pkgs/development/python-modules/nvidia-ml-py/default.nix @@ -40,10 +40,20 @@ buildPythonPackage rec { passthru.tests.tester-nvmlInit = cudaPackages.writeGpuTestPython { libraries = [ nvidia-ml-py ]; } '' - import pynvml - from pynvml.smi import nvidia_smi # noqa: F401 + from pynvml import ( + nvmlInit, + nvmlSystemGetDriverVersion, + nvmlDeviceGetCount, + nvmlDeviceGetHandleByIndex, + nvmlDeviceGetName, + ) - print(f"{pynvml.nvmlInit()=}") + nvmlInit() + print(f"Driver Version: {nvmlSystemGetDriverVersion()}") + + for i in range(nvmlDeviceGetCount()): + handle = nvmlDeviceGetHandleByIndex(i) + print(f"Device {i} : {nvmlDeviceGetName(handle)}") ''; meta = {