diff --git a/ucp/__init__.py b/ucp/__init__.py index 2edd17dd1..7486492c0 100644 --- a/ucp/__init__.py +++ b/ucp/__init__.py @@ -64,9 +64,14 @@ def _is_mig_device(handle): if _is_mig_device(handle): continue - total_memory = pynvml.nvmlDeviceGetMemoryInfo(handle).total - bar1_total = pynvml.nvmlDeviceGetBAR1MemoryInfo(handle).bar1Total + try: + bar1_total = pynvml.nvmlDeviceGetBAR1MemoryInfo(handle).bar1Total + except pynvml.nvml.NVMLError_NotSupported: + # Bar1 access not supported on this device, set it to + # zero (always lower than device memory). + bar1_total = 0 + total_memory = pynvml.nvmlDeviceGetMemoryInfo(handle).total if total_memory <= bar1_total: large_bar1[dev_idx] = True