Mask2Former broadcasting issue when running inference on model traced with GPU device
Author: matteot11
Pull request: no

### System Info

```
- System information: x86_64 GNU/Linux
- Ubuntu version: 18.04
- Python version: 3.8.12
- CUDA version: 11.1
- PyTorch version: 2.0.1
- transformers version: 4.31.0
```

### Who can help?

@amyeroberts
@sgugger 
@muellerzr


### Information

- [ ] The official example scripts
- [ ] My own modified scripts

### Tasks

- [ ] An officially supported task in the `examples` folder (such as GLUE/SQuAD, ...)
- [ ] My own task or dataset (give details below)

### Reproduction

```
import torch
from transformers import Mask2FormerForUniversalSegmentation

device = torch.device("cuda")
model = Mask2FormerForUniversalSegmentation.from_pretrained(
                "facebook/mask2former-swin-tiny-coco-instance",
                torchscript=True
                ).eval().to(device)
dummy_input = torch.randn((1,3,640,640)).to(device)
traced_model = torch.jit.trace(model, dummy_input)

with torch.no_grad():
    out = traced_model(torch.randn((2,3,640,640)).to(device))
    out = traced_model(torch.randn((2,3,640,640)).to(device))
```

The above code generates the following error when calling the **second** forward of `traced_model` (last line):
```
Traceback (most recent call last):
  File "mask2former_trace.py", line 14, in <module>
    out = traced_model(torch.randn((2,3,640,640)).to(device))
  File "~/python3.8/site-packages/torch/nn/modules/module.py", line 1501, in _call_impl
    return forward_call(*args, **kwargs)
RuntimeError: The following operation failed in the TorchScript interpreter.
Traceback of TorchScript (most recent call last):
~/python3.8/site-packages/torch/functional.py(378): einsum
~/python3.8/site-packages/transformers/models/mask2former/modeling_mask2former.py(2015): forward
~/python3.8/site-packages/torch/nn/modules/module.py(1488): _slow_forward
~/python3.8/site-packages/torch/nn/modules/module.py(1501): _call_impl
~/python3.8/site-packages/transformers/models/mask2former/modeling_mask2former.py(1852): forward
~/python3.8/site-packages/torch/nn/modules/module.py(1488): _slow_forward
~/python3.8/site-packages/torch/nn/modules/module.py(1501): _call_impl
~/python3.8/site-packages/transformers/models/mask2former/modeling_mask2former.py(2080): forward
~/python3.8/site-packages/torch/nn/modules/module.py(1488): _slow_forward
~/python3.8/site-packages/torch/nn/modules/module.py(1501): _call_impl
~/python3.8/site-packages/transformers/models/mask2former/modeling_mask2former.py(2271): forward
~/python3.8/site-packages/torch/nn/modules/module.py(1488): _slow_forward
~/python3.8/site-packages/torch/nn/modules/module.py(1501): _call_impl
~/python3.8/site-packages/transformers/models/mask2former/modeling_mask2former.py(2496): forward
~/python3.8/site-packages/torch/nn/modules/module.py(1488): _slow_forward
~/python3.8/site-packages/torch/nn/modules/module.py(1501): _call_impl
~/python3.8/site-packages/torch/jit/_trace.py(1056): trace_module
~/python3.8/site-packages/torch/jit/_trace.py(794): trace
mask2former_trace.py(10): <module>
RuntimeError: einsum(): subscript b has size 2 for operand 1 which does not broadcast with previously seen size 400
```

If I trace the model with batch size 2, i.e. `dummy_input = torch.randn((2,3,640,640)).to(device)`, the same error arises at the **first** forward call of `traced_model`

The issue seems to be [here](https://github.com/huggingface/transformers/blob/e42587f596181396e1c4b63660abf0c736b10dae/src/transformers/models/mask2former/modeling_mask2former.py#L2015)


### Expected behavior

When tracing on CPU, i.e. in the code above:
```
device = torch.device("cpu")
```
everything works fine. I would expect similar behaviour when tracing on GPU device.

**Additional notes**: 
I already tried tracing the model on CPU device, then moving `traced_model` (as well as the input tensors) to GPU, and running inference, but I got the following error:
```
RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cpu!
```
I know this is a known issue:
https://github.com/huggingface/transformers/issues/5664
https://github.com/huggingface/transformers/issues/22038
so I guess there should be some tensors in Mask2Former created at forward time with the same device as the input, and torchscript does not change that device when running on GPU. 
This is the reason why I need to trace the model on GPU.

Comments:
Hi @matteot11, thanks for reporting this and for providing such a detailed and clean issue report ❤️ 

Looking into it 🔍 

@matteot11 I'm going to open up a PR soon to resolve this and remove the einsum operations. In the meantime, if you need to be able to run a compiled model now, it will run on torch nightly (with a bunch of tracer warnings). 

Hi @amyeroberts, thanks for your fast reply.
With torch nightly I am able to correctly forward the `traced_model` multiple times (even if it was exported using `torch==2.0.1`). Thanks for the hint!

I don't know if this is expected, but when running the model traced on GPU, the following assert sometimes fails:
```
device = torch.device("cuda")
dummy_input = torch.randn((2,3,640,640)).to(device)
assert torch.isclose(model(dummy_input)[0], traced_model(dummy_input)[0]).all()
```
This does not happen when exporting the model to the CPU.

Waiting for your PR!

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.