Cannot export Deberta to TorchScript

[SohamTamba]

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.

[sgugger]

Yes, this model is not compatible with torchscript, cc @ArthurZucker

[ArthurZucker]

Thanks, will take that into account when refactoring

[SohamTamba]

Go away stalebot

[jinuhwang]

Any update here?

[ArthurZucker]

Just started working on this! 😉

[ArthurZucker]

Sorry! Seem like I had to postpone this! If anyone want to take over feel free to do it, otherwise will be my priority once #23909 is merge!

[ArthurZucker]

More delays given the recent sprints! But I think it should calm down during this summer! 😉

[dustinaxman]

Any update on this?

[LysandreJik]

I'll add a good first issue tag on this! Slightly more involved than a documentation change, but should be possible by someone willing to put a cycle into it.

[dustinaxman]

Thanks a ton!

[riyasachdeva04]

i would like to work on this

[LysandreJik]

Go ahead @riyasachdeva04, feel free to open a PR

[amannagarkar]

Can I take a crack at it too? @LysandreJik

[Szustarol]

Hey, I have noticed this issue has been silent for a while. I would like to start contributing, can I have a look at it? @LysandreJik @ArthurZucker

[LysandreJik]

First come first served, please feel free to open a PR and link this issue and we'll be happy to review!

[dustinaxman]

How is everyone doing with this? Is it worth it for me to give it a shot?