浏览代码

fix linux/amd gpu memory, convert to MB

Alex Cheema 10 月之前
父节点
当前提交
c819f675cb
共有 1 个文件被更改,包括 2 次插入2 次删除
  1. 2 2
      exo/topology/device_capabilities.py

+ 2 - 2
exo/topology/device_capabilities.py

@@ -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)