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}")