ONNX export fails for small image size

[gutcBrEl]

Describe the bug
I'm using the SwinTransformer for my application with fixed image size of 192x224 but dynamic batch size. While the training and inference using pytorch works well, the export to onnx fails with the following error:

minus_one_pos != -1 INTERNAL ASSERT FAILED at "/opt/pytorch/pytorch/torch/csrc/jit/passes/onnx/shape_type_inference.cpp":535, please report a bug to PyTorch. There are no examples for shape_has_zero = true && minus_one_pos == -1.

The error somehow will be produced in this line:

Swin-Transformer-Object-Detection/mmdet/models/backbones/swin_transformer.py

Line 142 in c7b2011

attn = attn.view(B_ // nW, nW, self.num_heads, N, N) + mask.unsqueeze(1).unsqueeze(0)

The error seems to occour for the given shapes:

mask shape: torch.Size([1, 49, 49])
attn shape: torch.Size([2, 24, 49, 49])
x shape: torch.Size([2, 49, 768])
num_heads: 24

Reproduction
The following script can be used to reproduce the bug:

import torch
from swin import SwinTransformer

model = SwinTransformer(
            depths=(2, 2, 6, 2),
            drop_path_rate=0,
            embed_dim=96,
            num_heads=(3, 6, 12, 24),
            out_indices=(0, 1, 2, 3),
            window_size=7,
        )

x = torch.rand([2, 3, 192, 224])

# run inference
with torch.no_grad():
    output = model(x)

# Export the model as onnx
onnx_path = "dummy_model.onnx"
torch.onnx.export(
    model,
    x,
    onnx_path,
    export_params=True,
    opset_version=17,
    do_constant_folding=True,
    input_names=['x'],
    output_names=['output'],
    dynamic_axes={'x': {0: 'batch_size'}, 'output': {0: 'batch_size'}}
)

Environment

torch                         2.2.0a0+6a974be
torch-tensorrt                2.2.0a0
torchdata                     0.7.0a0
torchtext                     0.16.0a0
torchvision                   0.17.0a0
onnx                          1.14.1

Error traceback
If applicable, paste the error trackback here.

Exception has occurred: RuntimeError
minus_one_pos != -1 INTERNAL ASSERT FAILED at "/opt/pytorch/pytorch/torch/csrc/jit/passes/onnx/shape_type_inference.cpp":535, please report a bug to PyTorch. There are no examples for shape_has_zero = true && minus_one_pos == -1.
  File "/workspace/dummy_export.py", line 24, in <module>
    torch.onnx.export(
RuntimeError: minus_one_pos != -1 INTERNAL ASSERT FAILED at "/opt/pytorch/pytorch/torch/csrc/jit/passes/onnx/shape_type_inference.cpp":535, please report a bug to PyTorch. There are no examples for shape_has_zero = true && minus_one_pos == -1.

Bug fix
The following work arounds avoids the error, but are not suitable for my application:

• increase image size to e.g. 384x480
• remove last stage by using only the first 3 values of the tuples depths, num_heads, out_indices
• use window_size 5