fix Project-MONAI#7265

Signed-off-by: KumoLiu <yunl@nvidia.com>
KumoLiu · Nov 29, 2023 · 31d4bee · 31d4bee
1 parent 860b7c3
commit 31d4bee
Showing 1 changed file with 21 additions and 18 deletions.
diff --git a/tests/test_swin_unetr.py b/tests/test_swin_unetr.py
@@ -24,7 +24,7 @@
 from monai.networks.nets.swin_unetr import PatchMerging, PatchMergingV2, SwinUNETR, filter_swinunetr
 from monai.networks.utils import copy_model_state
 from monai.utils import optional_import
-from tests.utils import assert_allclose, skip_if_downloading_fails, skip_if_no_cuda, skip_if_quick, testing_data_config
+from tests.utils import assert_allclose, skip_if_downloading_fails, skip_if_no_cuda, skip_if_quick, testing_data_config, SkipIfBeforePyTorchVersion
 
 einops, has_einops = optional_import("einops")
 
@@ -38,23 +38,25 @@
                 for img_size in ((64, 32, 192), (96, 32)):
                     for feature_size in [12]:
                         for norm_name in ["instance"]:
-                            test_case = [
-                                {
-                                    "spatial_dims": len(img_size),
-                                    "in_channels": in_channels,
-                                    "out_channels": out_channels,
-                                    "img_size": img_size,
-                                    "feature_size": feature_size,
-                                    "depths": depth,
-                                    "norm_name": norm_name,
-                                    "attn_drop_rate": attn_drop_rate,
-                                    "downsample": test_merging_mode[case_idx % 4],
-                                },
-                                (2, in_channels, *img_size),
-                                (2, out_channels, *img_size),
-                            ]
-                            case_idx += 1
-                            TEST_CASE_SWIN_UNETR.append(test_case)
+                            for use_checkpoint in [True, False]:
+                                test_case = [
+                                    {
+                                        "spatial_dims": len(img_size),
+                                        "in_channels": in_channels,
+                                        "out_channels": out_channels,
+                                        "img_size": img_size,
+                                        "feature_size": feature_size,
+                                        "depths": depth,
+                                        "norm_name": norm_name,
+                                        "attn_drop_rate": attn_drop_rate,
+                                        "downsample": test_merging_mode[case_idx % 4],
+                                        "use_checkpoint": use_checkpoint,
+                                    },
+                                    (2, in_channels, *img_size),
+                                    (2, out_channels, *img_size),
+                                ]
+                                case_idx += 1
+                                TEST_CASE_SWIN_UNETR.append(test_case)
 
 TEST_CASE_FILTER = [
     [
@@ -65,6 +67,7 @@
 ]
 
 
+@SkipIfBeforePyTorchVersion((1, 10))
 class TestSWINUNETR(unittest.TestCase):
     @parameterized.expand(TEST_CASE_SWIN_UNETR)
     @skipUnless(has_einops, "Requires einops")