Fix gpu parsing logic when numeric

mlcommons · Dec 8, 2023 · bd47fd0 · bd47fd0
1 parent b2aea76
commit bd47fd0
Showing 1 changed file with 1 addition and 3 deletions.
diff --git a/runners/mlcube_docker/mlcube_docker/docker_run.py b/runners/mlcube_docker/mlcube_docker/docker_run.py
@@ -275,12 +275,10 @@ def run(self) -> None:
                 cuda_visible_devices = cuda_visible_devices.replace("device=", "")
         else:
             cuda_visible_devices = num_gpus
-            if num_gpus == 0:
-                cuda_visible_devices = ""
 
         if cuda_visible_devices.isnumeric():
             cuda_visible_devices = str(list(range(cuda_visible_devices)))
-            cuda_visible_devices = cuda_visible_devices.replace(" ", "").replace("[","").replace("]","")
+            cuda_visible_devices = cuda_visible_devices.replace(" ", "")[1:-1]
 
         run_args += f" --env CUDA_VISIBLE_DEVICES={cuda_visible_devices}"