Fixed the metric so it works with ROCM librariers. Megvii-BaseDetection#1759

Author: rpehkone

yolox/utils/metric.py

@@ -26,10 +26,20 @@ def get_total_and_free_memory_in_Mb(cuda_device):
         "nvidia-smi --query-gpu=memory.total,memory.used --format=csv,nounits,noheader"
     )
     devices_info = devices_info_str.read().strip().split("\n")
-    if "CUDA_VISIBLE_DEVICES" in os.environ:
-        visible_devices = os.environ["CUDA_VISIBLE_DEVICES"].split(',')
-        cuda_device = int(visible_devices[cuda_device])
-    total, used = devices_info[int(cuda_device)].split(",")
+    if len(devices_info) > 1:
+        if "CUDA_VISIBLE_DEVICES" in os.environ:
+            visible_devices = os.environ["CUDA_VISIBLE_DEVICES"].split(',')
+            cuda_device = int(visible_devices[cuda_device])
+        total, used = devices_info[int(cuda_device)].split(",")
+    else:
+        devices_info_str = os.popen(
+            "rocm-smi --showmeminfo vram --csv"
+        )
+        devices_info = devices_info_str.read().strip().split("\n")
+        _, total, used = devices_info[1].split(',')
+        total = int(total) // (1024 * 1024)
+        used = int(used) // (1024 * 1024)
+
     return int(total), int(used)