fix linux/amd gpu memory, convert to MB

This commit is contained in:
Alex Cheema
2024-07-16 18:38:57 -07:00
parent e93a319ca4
commit c819f675cb
+2 -2
View File
@@ -52,9 +52,9 @@ def linux_device_capabilities() -> DeviceCapabilities:
print(f"NVIDIA device {gpu_name=} {gpu_memory_info=}")
return DeviceCapabilities(model=f"Linux Box ({gpu_name})", chip=gpu_name, memory=gpu_memory_info.total)
return DeviceCapabilities(model=f"Linux Box ({gpu_name})", chip=gpu_name, memory=gpu_memory_info.total // 2**20)
elif Device.DEFAULT == "AMD":
# TODO AMD support
return DeviceCapabilities(model="Linux Box (AMD)", chip="Unknown AMD", memory=psutil.virtual_memory().total)
return DeviceCapabilities(model="Linux Box (AMD)", chip="Unknown AMD", memory=psutil.virtual_memory().total // 2**20)
else:
return DeviceCapabilities(model=f"Linux Box (Device: {Device.DEFAULT})", chip=f"Unknown Chip (Device: {Device.DEFAULT})", memory=psutil.virtual_memory().total // 2**20)