Repository: transformers
Issue: Unable to use torch scripting to export Mask2Former model
Labels: Good Difficult Issue

### System Info

- `transformers` version: 4.34.1
- Platform: Linux-5.4.0-100-generic-x86_64-with-glibc2.17
- Python version: 3.8.18
- Huggingface_hub version: 0.17.3
- Safetensors version: 0.4.0
- Accelerate version: 0.23.0
- Accelerate config:    not found
- PyTorch version (GPU?): 2.1.0+cu121 (True)
- Tensorflow version (GPU?): not installed (NA)
- Flax version (CPU?/GPU?/TPU?): not installed (NA)
- Jax version: not installed
- JaxLib version: not installed
- Using GPU in script?: Yes, a GTX 1080 with 8 GB of VRAM
- Using distributed or parallel set-up in script?: No.

### Who can help?

@amyeroberts 

### 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

I am attempting to export the [Mask2Former model available in huggingface](https://huggingface.co/facebook/mask2former-swin-base-coco-panoptic) through [`torch.jit.script`](https://pytorch.org/docs/stable/generated/torch.jit.script.html). Here's a minimal reproducible example:

```python
import torch
from transformers import Mask2FormerForUniversalSegmentation
device = "cuda" if torch.cuda.is_available() else "cpu"
model = Mask2FormerForUniversalSegmentation.from_pretrained(
    "facebook/mask2former-swin-base-coco-panoptic", torchscript=True
    ).to(device)

scripted_model = torch.jit.script(model)
torch.jit.save(scripted_model, 'mask2former.pt')
```

By doing this, I get the following error using torch scripting (path to the offending file has been obfuscated for brevity):

```
torch.jit.frontend.NotSupportedError: Comprehension ifs are not supported yet:
  File "/home/.../huggingface/lib/python3.8/site-packages/transformers/models/mask2former/modeling_mask2former.py", line 2559
        if not return_dict:
            output = tuple(v for v in output.values() if v is not None)
            if loss is not None:
                output = ((loss)) + output
```

As a hack, I've changed my local installation so that comprehension ifs are removed:

```
    if not return_dict:
        outputs = []
        for v in output.values():
            if v is not None:
                outputs.append(v)
        output = tuple(outputs)    
```

This also occurs at line 2306 in the same file, so I've made the same changes there. Once I fix this, there is an error in the forward method for the SWIN backbone:

```
RuntimeError: 
'Optional[Tensor]' object has no attribute or method 'shape'.:
  File "/home/.../anaconda3/envs/huggingface/lib/python3.8/site-packages/transformers/models/swin/modeling_swin.py", line 313
    def forward(self, pixel_values: Optional[torch.FloatTensor]) -> Tuple[torch.Tensor, Tuple[int]]:
        _, num_channels, height, width = pixel_values.shape
                                         ~~~~~~~~~~~~~~~~~~ <--- HERE
        if num_channels != self.num_channels:
            raise ValueError(
```

The forward method for the SWIN backbone is confusing, as the input type is declared to be `Optional` but the output type is not. The definition of this method clearly indicates that a concrete tuple is to be returned.

As a final experiment, I've removed the `Optional` type declaration and tried to export it one more time:

```
aten::pad(Tensor self, SymInt[] pad, str mode="constant", float? value=None) -> Tensor:
Expected a value of type 'List[int]' for argument 'pad' but instead found type 'Tuple[int, Tensor]'.
:
  File "/home/.../anaconda3/envs/huggingface/lib/python3.8/site-packages/transformers/models/swin/modeling_swin.py", line 306
        if width % self.patch_size[1] != 0:
            pad_values = (0, self.patch_size[1] - width % self.patch_size[1])
            pixel_values = nn.functional.pad(pixel_values, pad_values)
                           ~~~~~~~~~~~~~~~~~ <--- HERE
        if height % self.patch_size[0] != 0:
            pad_values = (0, 0, 0, self.patch_size[0] - height % self.patch_size[0])
'SwinPatchEmbeddings.maybe_pad' is being compiled since it was called from 'SwinPatchEmbeddings.forward'
  File "/home/.../anaconda3/envs/huggingface/lib/python3.8/site-packages/transformers/models/swin/modeling_swin.py", line 319
            )
        # pad the input to be divisible by self.patch_size, if needed
        pixel_values = self.maybe_pad(pixel_values, height, width)
        ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ <--- HERE
        embeddings = self.projection(pixel_values)
        _, _, height, width = embeddings.shape
```

It seems that what is being put into the forward pass is not, in fact, a `torch.Tensor` when being scripted.

Is torch scripting this model not supported at this time or am I missing something?

### Expected behavior

The model successfully being exported to disk with torch scripting.