fix linux/amd gpu memory, convert to MB

c819f675cb468a6fe50be6adfd7fe118d8f52c18 · 2024-07-16 18:38:57 -0700 · Alex Cheema

Files touched

M exo/topology/device_capabilities.py

Diff

commit c819f675cb468a6fe50be6adfd7fe118d8f52c18
Author: Alex Cheema <alexcheema123@gmail.com>
Date:   Tue Jul 16 18:38:57 2024 -0700

    fix linux/amd gpu memory, convert to MB
---
 exo/topology/device_capabilities.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/exo/topology/device_capabilities.py b/exo/topology/device_capabilities.py
index fa74579f..57f97503 100644
--- a/exo/topology/device_capabilities.py
+++ b/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)

← e93a319c typo · back to Exo · Clarify readme iOS 7cb1ba55 →