The Conv2D layer in PatchMerging causes error

[MintCoffeeCat]

In the forward function of PatchMerging class, the msg_token and x was sent into self.reduction, which is a Conv2D layer with Conv2D(dim,2*dim,3,2,1) initialization.

        msg_token = window_reverse(x[:, :, 0].unsqueeze(2), 1, shuf_size, H//win_size, W//win_size, nchw=True)
        msg_token = self.reduction(msg_token).permute(0, 2, 3, 1)
        msg_token = self.norm(msg_token)
        if msg_token.shape[1] >= self.nxt_shuf_size:
            msg_token = window_partition(msg_token, 1, self.nxt_shuf_size)

        x = window_reverse(x[:, :, 1:], win_size, shuf_size, H, W, nchw=True)
        x = self.reduction(x).permute(0, 2, 3, 1)
        x = self.norm(x)

However, the later one sent into it will cause an error:

Exception has occurred: IndexError
tuple index out of range
File "/home/yun/Coding/MSG-Transformer/models/msg_transformer.py", line 298, in forward
x = self.reduction(x).permute(0, 2, 3, 1)
File "/home/yun/Coding/MSG-Transformer/models/msg_transformer.py", line 384, in forward
x = self.downsample(x)
File "/home/yun/Coding/MSG-Transformer/models/msg_transformer.py", line 579, in forward_features
x = layer(x)
File "/home/yun/Coding/MSG-Transformer/models/msg_transformer.py", line 585, in forward
x = self.forward_features(x)
File "/home/yun/Coding/MSG-Transformer/main.py", line 168, in train_one_epoch
outputs = model(samples)
File "/home/yun/Coding/MSG-Transformer/main.py", line 134, in main
train_one_epoch(config, model, criterion, data_loader_train, optimizer, epoch, mixup_fn, lr_scheduler)
File "/home/yun/Coding/MSG-Transformer/main.py", line 356, in
main(config)

I tried change the order, letting x sent to Conv2D first and msg_token later but also cause this error.
Some of my environments here:
python 3.8.15
torchaudio 0.12.0 py38_cu113 pytorch
torchvision 0.13.0 py38_cu113 pytorch
pytorch 1.12.0 py3.8_cuda11.3_cudnn8.3.2_0 pytorch
timm 0.4.12 pypi_0 pypi

Would you please give some suggestions?