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:
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
|
||||
|
||||
Reference in New Issue
Block a user