fix

codalab · Aug 16, 2023 · b43fd07 · b43fd07
1 parent 184fd95
commit b43fd07
Showing 1 changed file with 5 additions and 3 deletions.
diff --git a/codalab/worker/docker_utils.py b/codalab/worker/docker_utils.py
@@ -164,9 +164,11 @@ def get_nvidia_devices(self, use_docker=True):
                 stdout=True,
                 remove=True,
             )
-            gpus = output.decode()
-            if gpus.find("==========\n== CUDA ==\n==========") != -1:  # need to remove header
-                gpus = gpus.split("\n")[15:-1]
+            gpu_info = output.decode()
+            if gpu_info.find("==========") != -1:  # need to remove header
+                gpus = gpu_info.split("\n")[15:-1]
+            else:
+                gpus = gpu_info.split("\n")[:-1]
 
         else:
             # use the singularity runtime to run nvidia-smi