fix(utils): occupy_mem based on CUDA_VISIBLE_DEVICES (#1139)

fix(utils): occupy_mem based on CUDA_VISIBLE_DEVICES
Megvii-BaseDetection · Feb 22, 2022 · d98da01 · d98da01
1 parent d99bee0
commit d98da01
Showing 1 changed file with 3 additions and 0 deletions.
diff --git a/yolox/utils/metric.py b/yolox/utils/metric.py
@@ -24,6 +24,9 @@ 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(",")
     return int(total), int(used)