workaround for Project-MONAI#7575

Signed-off-by: YunLiu <55491388+KumoLiu@users.noreply.github.com>
KumoLiu · Mar 25, 2024 · 9319a3c · 9319a3c
1 parent b87ae5e
commit 9319a3c
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 5 deletions.
diff --git 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
@@ -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; '