diff --git a/monai/apps/nnunet/nnunetv2_runner.py b/monai/apps/nnunet/nnunetv2_runner.py index 0308bcb1d6..25038238ac 100644 --- a/monai/apps/nnunet/nnunetv2_runner.py +++ b/monai/apps/nnunet/nnunetv2_runner.py @@ -533,15 +533,15 @@ def train_single_model_command(self, config, fold, gpu_id, kwargs): gpu_ids_str = "" for _i in range(len(gpu_id)): gpu_ids_str += f"{gpu_id[_i]}," - set_visible_device = f"CUDA_VISIBLE_DEVICES={gpu_ids_str[:-1]}" + device_setting = f"CUDA_VISIBLE_DEVICES={gpu_ids_str[:-1]}" else: - set_visible_device = f"CUDA_VISIBLE_DEVICES={gpu_id[0]}" + device_setting = f"CUDA_VISIBLE_DEVICES={gpu_id[0]}" else: - set_visible_device = f"CUDA_VISIBLE_DEVICES={gpu_id}" + device_setting = f"CUDA_VISIBLE_DEVICES={gpu_id}" num_gpus = 1 if isinstance(gpu_id, int) or len(gpu_id) == 1 else len(gpu_id) cmd = ( - f"{set_visible_device} nnUNetv2_train " + f"{device_setting} nnUNetv2_train " + f"{self.dataset_name_or_id} {config} {fold} " + f"-tr {self.trainer_class_name} -num_gpus {num_gpus}" ) diff --git a/tests/test_set_visible_devices.py b/tests/test_set_visible_devices.py index 7860656b3d..b4f44957a2 100644 --- a/tests/test_set_visible_devices.py +++ b/tests/test_set_visible_devices.py @@ -14,7 +14,7 @@ import os import unittest -from tests.utils import skip_if_no_cuda +from tests.utils import SkipIfAtLeastPyTorchVersion, skip_if_no_cuda class TestVisibleDevices(unittest.TestCase): @@ -25,6 +25,7 @@ def run_process_and_get_exit_code(code_to_execute): return int(bin(value).replace("0b", "").rjust(16, "0")[:8], 2) @skip_if_no_cuda + @SkipIfAtLeastPyTorchVersion((2, 2, 1)) def test_visible_devices(self): num_gpus_before = self.run_process_and_get_exit_code( 'python -c "import os; import torch; '