Loading and running on CPU, the RoBERTa model traced/saved on GPU.
Author: HamidShojanazeri
Pull request: no
Labels: name: wontfix

# ❓ Questions & Help


## Details
I have installed the transformers from source and I am trying to trace/save roberta_base model for sequence classification on GPU and load it on CPU, similar to the issue [#5664](https://github.com/huggingface/transformers/issues/5664) that was addressed [PR #5773](https://github.com/huggingface/transformers/pull/5773)  and I am facing similar issues with position_embeddings. I was wondering is any step is missing or RoBERTa is also covered in [PR #5773](https://github.com/huggingface/transformers/pull/5773). Thanks.

## Sample Script

```python
import transformers
from pathlib import Path
import os
import json
import torch
from transformers import (AutoModelForSequenceClassification, AutoTokenizer, AutoModelForQuestionAnswering,
 AutoModelForTokenClassification, AutoConfig)

print('Transformers version',transformers.__version__)
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
dummy_input = "This is a dummy input for torch jit trace"
max_length=20
config = AutoConfig.from_pretrained('roberta-base',num_labels=2,torchscript=True)
model = AutoModelForSequenceClassification.from_pretrained('roberta-base', config=config)
tokenizer = AutoTokenizer.from_pretrained('roberta-base',do_lower_case=True)
inputs = tokenizer.encode_plus(dummy_input,max_length = max_length,pad_to_max_length = True,truncation=True, add_special_tokens = True, return_tensors = 'pt')

print(inputs.keys())
input_ids = inputs["input_ids"].to(device)
attention_mask = inputs["attention_mask"].to(device)
model.to(device).eval()
traced_model = torch.jit.trace(model, (input_ids,attention_mask))
torch.jit.save(traced_model, "Roberta_cuda.pt")

print(traced_model.graph)
print("\n")
print("Load model onto CPU")
loaded = torch.jit.load("Roberta_cuda.pt", map_location=torch.device("cpu"))
print("\n")
print(loaded.graph)
outputs = loaded(input_ids.to("cpu"),attention_mask.to("cpu"))
print(outputs)
```

## Error 

Traceback (most recent call last):
  File "test.py", line 100, in <module>
    outputs = loaded(input_ids.to("cpu"),attention_mask.to("cpu"))
  File "/home/ubuntu/anaconda3/envs/pytorch_latest_p36/lib/python3.6/site-packages/torch/nn/modules/module.py", line 550, in __call__
    result = self.forward(*input, **kwargs)
RuntimeError: The following operation failed in the TorchScript interpreter.
Traceback of TorchScript, serialized code (most recent call last):
  File "code/__torch__/transformers/modeling_roberta.py", line 10, in forward
    attention_mask: Tensor) -> Tuple[Tensor]:
    _0 = self.classifier
    _1 = (self.roberta).forward(input_ids, attention_mask, )
          ~~~~~~~~~~~~~~~~~~~~~ <--- HERE
    return ((_0).forward(_1, ),)
class RobertaModel(Module):
  File "code/__torch__/transformers/modeling_roberta.py", line 32, in forward
    _9 = torch.to(extended_attention_mask, 6, False, False, None)
    attention_mask0 = torch.mul(torch.rsub(_9, 1., 1), CONSTANTS.c0)
    _10 = (_3).forward((_4).forward(input_ids, input, ), attention_mask0, )
                        ~~~~~~~~~~~ <--- HERE
    _11 = (_2).forward(_10, )
    return _10
  File "code/__torch__/transformers/modeling_roberta.py", line 59, in forward
    input0 = torch.to(_19, dtype=4, layout=0, device=torch.device("cuda:0"), pin_memory=False, non_blocking=False, copy=False, memory_format=None)
    _20 = (_16).forward(input_ids, )
    _21 = (_15).forward(input0, )
           ~~~~~~~~~~~~ <--- HERE
    _22 = (_14).forward(input, )
    input1 = torch.add(torch.add(_20, _21, alpha=1), _22, alpha=1)
  File "code/__torch__/torch/nn/modules/sparse/___torch_mangle_0.py", line 7, in forward
  def forward(self: __torch__.torch.nn.modules.sparse.___torch_mangle_0.Embedding,
    input: Tensor) -> Tensor:
    position_embeddings = torch.embedding(self.weight, input, 1, False, False)
                          ~~~~~~~~~~~~~~~ <--- HERE
    return position_embeddings

Traceback of TorchScript, original code (most recent call last):
/home/ubuntu/anaconda3/envs/pytorch_latest_p36/lib/python3.6/site-packages/torch/nn/functional.py(1724): embedding
/home/ubuntu/anaconda3/envs/pytorch_latest_p36/lib/python3.6/site-packages/torch/nn/modules/sparse.py(114): forward
/home/ubuntu/anaconda3/envs/pytorch_latest_p36/lib/python3.6/site-packages/torch/nn/modules/module.py(534): _slow_forward
/home/ubuntu/anaconda3/envs/pytorch_latest_p36/lib/python3.6/site-packages/torch/nn/modules/module.py(548): __call__
/home/ubuntu/anaconda3/envs/pytorch_latest_p36/lib/python3.6/site-packages/transformers/modeling_bert.py(202): forward
/home/ubuntu/anaconda3/envs/pytorch_latest_p36/lib/python3.6/site-packages/transformers/modeling_roberta.py(76): forward
/home/ubuntu/anaconda3/envs/pytorch_latest_p36/lib/python3.6/site-packages/torch/nn/modules/module.py(534): _slow_forward
/home/ubuntu/anaconda3/envs/pytorch_latest_p36/lib/python3.6/site-packages/torch/nn/modules/module.py(548): __call__
/home/ubuntu/anaconda3/envs/pytorch_latest_p36/lib/python3.6/site-packages/transformers/modeling_bert.py(792): forward
/home/ubuntu/anaconda3/envs/pytorch_latest_p36/lib/python3.6/site-packages/torch/nn/modules/module.py(534): _slow_forward
/home/ubuntu/anaconda3/envs/pytorch_latest_p36/lib/python3.6/site-packages/torch/nn/modules/module.py(548): __call__
/home/ubuntu/anaconda3/envs/pytorch_latest_p36/lib/python3.6/site-packages/transformers/modeling_roberta.py(344): forward
/home/ubuntu/anaconda3/envs/pytorch_latest_p36/lib/python3.6/site-packages/torch/nn/modules/module.py(534): _slow_forward
/home/ubuntu/anaconda3/envs/pytorch_latest_p36/lib/python3.6/site-packages/torch/nn/modules/module.py(548): __call__
/home/ubuntu/anaconda3/envs/pytorch_latest_p36/lib/python3.6/site-packages/torch/jit/__init__.py(1027): trace_module
/home/ubuntu/anaconda3/envs/pytorch_latest_p36/lib/python3.6/site-packages/torch/jit/__init__.py(875): trace
test.py(91): <module>
RuntimeError: Expected object of device type cuda but got device type cpu for argument #1 'self' in call to _th_index_select

Comments:
This issue has been automatically marked as stale because it has not had recent activity. It will be closed if no further activity occurs. Thank you for your contributions.


@mfuntowicz  I was wondering if there is any update/ work around or a PR in progress to solve this issue. Thanks.