Fix: fallback to torch API when NVML memory query is not supported (#23426)

Co-authored-by: ronnie_zheng <zl19940307@163.com>
This commit is contained in:
Jinghong Li
2026-04-23 19:26:04 +03:00
committed by GitHub
co-authored by ronnie_zheng
parent 86ed0680d7
commit c6872fc8fb
@@ -514,7 +514,10 @@ class NvmlCudaPlatform(CudaPlatformBase):
def get_device_total_memory(cls, device_id: int = 0) -> int:
physical_device_id = device_id_to_physical_device_id(device_id)
handle = pynvml.nvmlDeviceGetHandleByIndex(physical_device_id)
return int(pynvml.nvmlDeviceGetMemoryInfo(handle).total)
try:
return int(pynvml.nvmlDeviceGetMemoryInfo(handle).total)
except pynvml.NVMLError_NotSupported:
return int(torch.cuda.get_device_properties(device_id).total_memory)
@classmethod
@with_nvml_context