transformers: Cannot export Deberta to TorchScript

System Info

transformers-cli env

- `transformers` version: 4.10.2
- Platform: Linux-3.10.0-1127.18.2.el7.x86_64-x86_64-with-glibc2.23
- Python version: 3.9.13
- PyTorch version (GPU?): 1.9.0 (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?: No
- Using distributed or parallel set-up in script?: No

Who can help?

@sgugger

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

I am trying to convert the Deberta Model to TorchScript using the instructions provided in the HF tutorial.

Code:

import torch
from transformers import AutoTokenizer, AutoModel

tokenizer = AutoTokenizer.from_pretrained("microsoft/deberta-base")
model = AutoModel.from_pretrained("microsoft/deberta-base", torchscript=True)

tokenized_dict = tokenizer(
    ["Is this working",], ["Not yet",], 
    return_tensors="pt"
)
input_tuple = (tokenized_dict['input_ids'], tokenized_dict['attention_mask'])

traced_model = torch.jit.trace(model, input_tuple)
torch.jit.save(traced_model, "compiled_deberta.pt")

Error Message: From torch.jit.save:

Could not export Python function call 'XSoftmax'. Remove calls to Python functions before export. Did you forget to add @script or @script_method annotation? If this is a nn.ModuleList, add it to __constants__:

Expected behavior

The Traced model should be successfully saved. After loading, it should have the same functional behavior as the model it was traced from.