Swin trace is not correctedly for dynamic batch.
Author: MenglingD
Pull request: no
Labels: name: bug

### System Info

transformers: v4.21.2
system: centos
python version: 3.6

### Who can help?

@sgugger @NielsRogge

### Information

- [ ] The official example scripts
- [X] My own modified scripts

### Tasks

- [ ] An officially supported task in the `examples` folder (such as GLUE/SQuAD, ...)
- [X] My own task or dataset (give details below)

### Reproduction

```python3
import torch
from models import build_model
from config import _C, _update_config_from_file

config = _C.clone()
_update_config_from_file(config, "configs/swin/swin_tiny_patch4_window7_224.yaml")
model = build_model(config).cuda().eval()

input1 = torch.randn((1, 3, 224, 224)).to("cuda")
input2 = torch.randn((2, 3, 224, 224)).to("cuda")
jit_model = torch.jit.trace(model, input1)

assert((model(input1) - jit_model(input1)).abs().sum() == 0)
assert((model(input2) - jit_model(input2)).abs().sum() == 0)
```

### Expected behavior

Expected nothing, but got `AssertionError`.

Comments:
Wondering why this wasn't caught by the torchscript tests, cc @michaelbenayoun 

The dimension of batch is casted  to integer forcelly, which leads to swin cann't support dynamic batch [modeling_swin.py#L218](https://github.com/huggingface/transformers/blob/main/src/transformers/models/swin/modeling_swin.py#L218)：

```python3
def window_reverse(windows, window_size, height, width):
    """
    Merges windows to produce higher resolution features.
    """
    batch_size = math.floor(windows.shape[0] / (height * width / window_size / window_size))
    windows = windows.view(batch_size, height // window_size, width // window_size, window_size, window_size, -1)
    windows = windows.permute(0, 1, 3, 2, 4, 5).contiguous().view(batch_size, height, width, -1)
    return windows
```

And the following modifications has same semantics but supports dynamic batch:

```python3
def window_reverse(windows, window_size, height, width):
    """
    Merges windows to produce higher resolution features.
    """
    channels = int(windows.shape[-1])
    windows = windows.view(-1, height // window_size, width // window_size, window_size, window_size, channels)
    windows = windows.permute(0, 1, 3, 2, 4, 5).contiguous().view(-1, height, width, channels)
    return windows
```

ps: I have pulled request before #18807, but I can't `make fixup` as the python version is too old(python3.6.8), and I have close that PR. I may need to trouble you to fix it.

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the [contributing guidelines](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md) are likely to be ignored.

Hi @MenglingD,
I am not even able to trace the model except if I disable `SwinLayer.set_shift_and_window_size` [here](https://github.com/huggingface/transformers/blob/v4.21-release/src/transformers/models/swin/modeling_swin.py#L649).
But when disabling it, I am able to get it to work with the change you recommended.

I think we can apply this fix, but as long as `SwinLayer.set_shift_and_window_size` is not tracing friendly, it won't fix a thing.

> Hi @MenglingD, I am not even able to trace the model except if I disable `SwinLayer.set_shift_and_window_size` [here](https://github.com/huggingface/transformers/blob/v4.21-release/src/transformers/models/swin/modeling_swin.py#L649). But when disabling it, I am able to get it to work with the change you recommended.
> 
> I think we can apply this fix, but as long as `SwinLayer.set_shift_and_window_size` is not tracing friendly, it won't fix a thing.

OK, thanks.

I wander why that `SwinLayer.set_shift_and_window_size is not tracing friendly`. The `input_dimensions` is fixed at the deployment phase, and it is ok for jit.trace which will choose one of branch of if-else statement. I am interesting in this trace error, could you please provide more message for it?

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the [contributing guidelines](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md) are likely to be ignored.