From 953037f982cf06fdf6ecf5a8d9546ec01ecee67f Mon Sep 17 00:00:00 2001 From: rpehkone <-> Date: Mon, 15 Apr 2024 21:53:34 +0300 Subject: [PATCH] Fixed the metric so it works with ROCM librariers. #1759 --- yolox/utils/metric.py | 18 ++++++++++++++---- 1 file changed, 14 insertions(+), 4 deletions(-) diff --git a/yolox/utils/metric.py b/yolox/utils/metric.py index 506b58281..da17d44dd 100644 --- a/yolox/utils/metric.py +++ b/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)