Repository: Project-MONAI/MONAI

Issue:
PatchMerging padding error
**Describe the bug**
```py
import torch

from monai.networks.nets.swin_unetr import PatchMerging

dim=10
t = PatchMerging(dim)(torch.zeros((1, 21, 20, 20, dim)))
print(t.shape)
```

```
Traceback (most recent call last):
  File "my_test_3.py", line 7, in <module>
    t = PatchMerging(dim)(torch.zeros((1, 21, 20, 20, dim)))
  File "py38/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1130, in _call_impl
    return forward_call(*input, **kwargs)
  File "MONAI/monai/networks/nets/swin_unetr.py", line 701, in forward
    x = torch.cat([x0, x1, x2, x3, x4, x5, x6, x7], -1)
RuntimeError: Sizes of tensors must match except in dimension 4. Expected size 11 but got size 10 for tensor number 1 in the list.
```


Discussion:

Patch:
diff --git a/monai/networks/blocks/patchembedding.py b/monai/networks/blocks/patchembedding.py
--- a/monai/networks/blocks/patchembedding.py
+++ b/monai/networks/blocks/patchembedding.py
@@ -180,7 +180,7 @@ def forward(self, x):
                 x = F.pad(x, (0, 0, 0, 0, 0, self.patch_size[0] - d % self.patch_size[0]))
 
         elif len(x_shape) == 4:
-            _, _, h, w = x.size()
+            _, _, h, w = x_shape
             if w % self.patch_size[1] != 0:
                 x = F.pad(x, (0, self.patch_size[1] - w % self.patch_size[1]))
             if h % self.patch_size[0] != 0:
diff --git a/monai/networks/nets/swin_unetr.py b/monai/networks/nets/swin_unetr.py
--- a/monai/networks/nets/swin_unetr.py
+++ b/monai/networks/nets/swin_unetr.py
@@ -350,7 +350,7 @@ def window_reverse(windows, window_size, dims):
 
     elif len(dims) == 3:
         b, h, w = dims
-        x = windows.view(b, h // window_size[0], w // window_size[0], window_size[0], window_size[1], -1)
+        x = windows.view(b, h // window_size[0], w // window_size[1], window_size[0], window_size[1], -1)
         x = x.permute(0, 1, 3, 2, 4, 5).contiguous().view(b, h, w, -1)
     return x
 
@@ -570,8 +570,8 @@ def forward_part1(self, x, mask_matrix):
             b, h, w, c = x.shape
             window_size, shift_size = get_window_size((h, w), self.window_size, self.shift_size)
             pad_l = pad_t = 0
-            pad_r = (window_size[0] - h % window_size[0]) % window_size[0]
-            pad_b = (window_size[1] - w % window_size[1]) % window_size[1]
+            pad_b = (window_size[0] - h % window_size[0]) % window_size[0]
+            pad_r = (window_size[1] - w % window_size[1]) % window_size[1]
             x = F.pad(x, (0, 0, pad_l, pad_r, pad_t, pad_b))
             _, hp, wp, _ = x.shape
             dims = [b, hp, wp]
@@ -689,7 +689,7 @@ def forward(self, x):
             b, d, h, w, c = x_shape
             pad_input = (h % 2 == 1) or (w % 2 == 1) or (d % 2 == 1)
             if pad_input:
-                x = F.pad(x, (0, 0, 0, d % 2, 0, w % 2, 0, h % 2))
+                x = F.pad(x, (0, 0, 0, w % 2, 0, h % 2, 0, d % 2))
             x0 = x[:, 0::2, 0::2, 0::2, :]
             x1 = x[:, 1::2, 0::2, 0::2, :]
             x2 = x[:, 0::2, 1::2, 0::2, :]
@@ -782,7 +782,7 @@ def __init__(
         """
         Args:
             dim: number of feature channels.
-            depths: number of layers in each stage.
+            depth: number of layers in each stage.
             num_heads: number of attention heads.
             window_size: local window size.
             drop_path: stochastic depth rate.


Test patch:
diff --git a/tests/test_swin_unetr.py b/tests/test_swin_unetr.py
--- a/tests/test_swin_unetr.py
+++ b/tests/test_swin_unetr.py
@@ -16,7 +16,7 @@
 from parameterized import parameterized
 
 from monai.networks import eval_mode
-from monai.networks.nets.swin_unetr import SwinUNETR
+from monai.networks.nets.swin_unetr import PatchMerging, SwinUNETR
 from monai.utils import optional_import
 
 einops, has_einops = optional_import("einops")
@@ -26,26 +26,24 @@
     for in_channels in [1]:
         for depth in [[2, 1, 1, 1], [1, 2, 1, 1]]:
             for out_channels in [2]:
-                for img_size in [64]:
+                for img_size in ((64, 32, 192), (96, 32)):
                     for feature_size in [12]:
                         for norm_name in ["instance"]:
-                            for nd in (2, 3):
-                                test_case = [
-                                    {
-                                        "in_channels": in_channels,
-                                        "out_channels": out_channels,
-                                        "img_size": (img_size,) * nd,
-                                        "feature_size": feature_size,
-                                        "depths": depth,
-                                        "norm_name": norm_name,
-                                        "attn_drop_rate": attn_drop_rate,
-                                    },
-                                    (2, in_channels, *([img_size] * nd)),
-                                    (2, out_channels, *([img_size] * nd)),
-                                ]
-                                if nd == 2:
-                                    test_case[0]["spatial_dims"] = 2  # type: ignore
-                                TEST_CASE_SWIN_UNETR.append(test_case)
+                            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,
+                                },
+                                (2, in_channels, *img_size),
+                                (2, out_channels, *img_size),
+                            ]
+                            TEST_CASE_SWIN_UNETR.append(test_case)
 
 
 class TestSWINUNETR(unittest.TestCase):
@@ -84,6 +82,11 @@ def test_ill_arg(self):
                 drop_rate=0.4,
             )
 
+    def test_patch_merging(self):
+        dim = 10
+        t = PatchMerging(dim)(torch.zeros((1, 21, 20, 20, dim)))
+        self.assertEqual(t.shape, torch.Size([1, 11, 10, 10, 20]))
+
 
 if __name__ == "__main__":
     unittest.main()