GPTDoubleHeadsModel Unexpected node type: onnx:: Sub 
Author: mihail911
Pull request: no

# 🐛 Bug

## Information

Model I am using (Bert, XLNet ...): GPT2DoubleHeadsModel

Language I am using the model on (English, Chinese ...): English

The problem arises when using:
* [x] the official example scripts: (give details below)
* [ ] my own modified scripts: (give details below)

The tasks I am working on is:
* [ ] an official GLUE/SQUaD task: (give the name)
* [x] my own task or dataset: (give details below)

## To reproduce

Steps to reproduce the behavior:
I've been following the ipython notebook provided [here](https://github.com/huggingface/transformers/blob/master/notebooks/04-onnx-export.ipynb)

1. Take an off-the-shelf pretrained `gpt` model and export to onnx format using the following script: 
```
import torch
from transformers import (GPT2Config, GPT2Model, GPT2Tokenizer)

# use_cache is True by default in GPT2Model. Here we wrap a class to disable past state output.
class GPT2DoubleHeadsModelNoPastState(GPT2DoubleHeadsModel):
    def __init__(self, config):
        super().__init__(config)

    def forward(self, input_ids, token_type_ids):
        return super().forward(input_ids, past=None, attention_mask=None, token_type_ids=token_type_ids, use_cache=False)

model_name="gpt2"
config = GPT2Config.from_pretrained(model_name)
tokenizer = GPT2Tokenizer.from_pretrained(model_name)
model = GPT2ModelNoPastState.from_pretrained(model_name)

example_inputs = tokenizer.encode_plus("This is a sample input", return_tensors="pt")
del example_inputs["attention_mask"]
example_outputs = model(**example_inputs)

input_names = ['input_ids',  'token_type_ids']
output_names=['output_1', 'output_2']
dynamic_axes={'input_ids': {0: 'batch_size', 1: 'num_choices', 2: 'seq_len'}, 
     'token_type_ids': {0: 'batch_size', 1: 'num_choices', 2: 'seq_len'},
     'output_1': {0: 'batch_size', 1: 'num_choices': 2: 'seq_len', 3: 'vocab_size'}, 
     'output_2': {0: 'batch_size', 1: 'num_choices'}
}
output_path='gpt2.onnx'
torch.onnx.export(model=model,
                  args=(example_inputs[input_names[0]].unsqueeze(0), example_inputs[input_names[1]].unsqueeze(0)),
                  f=output_path,
                  input_names=input_names,
                  output_names=output_names,
                  example_outputs=example_outputs,
                  dynamic_axes=dynamic_axes,
                  do_constant_folding=True,
                  opset_version=11,
                  use_external_data_format=False)
```

This script is based off of #4805 

2.

After invoking the above, I get the error:
```
.../torch/onnx/symbolic_helper.py", line 87...
RuntimeError: Unexpected node type: onnx::Sub
```

## Expected behavior

I would expect this to work successfully, and unfortunately I'm not exactly sure how to interpret this error. There's not a lot of documentation online. 

## Environment info
<!-- You can run the command `transformers-cli env` and copy-and-paste its output below.
     Don't forget to fill out the missing fields in that output! -->
     
- `transformers` version: Commit 0e1869cc286d607f1598506be7bd1312b76ca82c
- Onnxruntime: 1.3.0
- Python version: 3.6.10
- PyTorch version (GPU?): 1.5.0+cu101
- Using GPU in script?: Yes

Thanks for your help! @mfuntowicz @tianleiwu 


Comments:
Hi @mihail911, thanks for reporting the issue and the script to reproduce.

I can confirm the issue, as it seems to happen on PyTorch side, I suspect it's a bug on their side. @tianleiwu should we forward the issue on PyTorch issue tracker? 

Slightly updated the script to avoid errors: 

```python
import torch
from transformers import (GPT2Config, GPT2Model, GPT2Tokenizer, GPT2DoubleHeadsModel)

# use_cache is True by default in GPT2Model. Here we wrap a class to disable past state output.
class GPT2DoubleHeadsModelNoPastState(GPT2DoubleHeadsModel):
    def __init__(self, config):
        super().__init__(config)

    def forward(self, input_ids, token_type_ids):
        return super().forward(input_ids, past=None, attention_mask=None, token_type_ids=token_type_ids, use_cache=False)

model_name="gpt2"
config = GPT2Config.from_pretrained(model_name)
tokenizer = GPT2Tokenizer.from_pretrained(model_name)
model = GPT2DoubleHeadsModelNoPastState.from_pretrained(model_name)

example_inputs = tokenizer.encode_plus("This is a sample input", return_tensors="pt")
del example_inputs["attention_mask"]
example_outputs = model(**example_inputs)

input_names = ['input_ids',  'token_type_ids']
output_names=['output_1', 'output_2']
dynamic_axes={'input_ids': {0: 'batch_size', 1: 'num_choices', 2: 'seq_len'}, 
     'token_type_ids': {0: 'batch_size', 1: 'num_choices', 2: 'seq_len'},
     'output_1': {0: 'batch_size', 1: 'num_choices', 2: 'seq_len', 3: 'vocab_size'}, 
     'output_2': {0: 'batch_size', 1: 'num_choices'}
}
output_path='gpt2.onnx'
torch.onnx.export(model=model,
                  args=(example_inputs[input_names[0]].unsqueeze(0), example_inputs[input_names[1]].unsqueeze(0)),
                  f=output_path,
                  input_names=input_names,
                  output_names=output_names,
                  example_outputs=example_outputs,
                  dynamic_axes=dynamic_axes,
                  do_constant_folding=True,
                  opset_version=11,
                  use_external_data_format=False)
```

@mfuntowicz, I've forwarded the issue to the developer of pytorch onnx exporter.

I did narrow down the issue to [one line](https://github.com/huggingface/transformers/blob/ca5e1cdf8e314288bd0242a531815a6c75d8178e/src/transformers/modeling_utils.py#L2056). A walk-around is to add int() to cast data type:

Before:
```
cls_index = torch.full_like(hidden_states[..., :1, :], hidden_states.shape[-2] - 1, dtype=torch.long,)
```

After:
```
cls_index = torch.full_like(hidden_states[..., :1, :], int(hidden_states.shape[-2]) - 1, dtype=torch.long,)
```
@mihail911, could you try this (need install transformers from source) to see whether you can export the model?





Yes I am able to export the model. Thanks @tianleiwu @mfuntowicz!