Fix validation when using --gpus flag

mlcommons · Nov 30, 2023 · da9d74a · da9d74a
1 parent 7a5fea8
commit da9d74a
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/runners/mlcube_docker/mlcube_docker/docker_run.py b/runners/mlcube_docker/mlcube_docker/docker_run.py
@@ -266,7 +266,8 @@ def run(self) -> None:
         if extra_args:
             run_args += " " + extra_args
 
-        cuda_visible_devices = self.mlcube.runner["--gpus"] if "--gpus" in self.mlcube.runner else num_gpus
+        valid_gpu_flag = "--gpus" in self.mlcube.runner and self.mlcube.runner["--gpus"] is not None
+        cuda_visible_devices = self.mlcube.runner["--gpus"] if valid_gpu_flag else num_gpus
         run_args += f" --env CUDA_VISIBLE_DEVICES={cuda_visible_devices}"
 
         if "entrypoint" in self.mlcube.tasks[self.task]: