Llama-2-13B / system /pynvml_check.py
jordigonzm's picture
check installation files
8d4d3e5
raw
history blame
No virus
705 Bytes
import pynvml
try:
pynvml.nvmlInit()
device_count = pynvml.nvmlDeviceGetCount()
print(f"Number of GPUs: {device_count}")
for i in range(device_count):
handle = pynvml.nvmlDeviceGetHandleByIndex(i)
name = pynvml.nvmlDeviceGetName(handle)
memory_info = pynvml.nvmlDeviceGetMemoryInfo(handle)
print(f"GPU {i}: {name.decode('utf-8')}")
print(f" Memory Total: {memory_info.total / (1024 ** 2)} MB")
print(f" Memory Free: {memory_info.free / (1024 ** 2)} MB")
print(f" Memory Used: {memory_info.used / (1024 ** 2)} MB")
pynvml.nvmlShutdown()
except pynvml.NVMLError as e:
print(f"Failed to initialize NVML: {e}")