Unable to use torch scripting to export Mask2Former model

### 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: outpu…