diff --git a/runners/mlcube_docker/mlcube_docker/docker_run.py b/runners/mlcube_docker/mlcube_docker/docker_run.py index 2c84140..b3b021f 100644 --- 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}"