Get "RuntimeError: forward() Expected a value of type 'Tensor' for argument 'input_ids' but instead found type 'list'" while loading Torchscript models Converting from Transformers

Beginners
1

I followed the tutorial here, modified the code for ALBERT and ran them on Google Colab:
Step 1

from transformers import AlbertModel, AlbertTokenizer, AlbertConfig
import torch
enc = AlbertTokenizer.from_pretrained("albert-xxlarge-v2")
text = "[CLS] Who was Jim Henson ? [SEP] Jim Henson was a puppeteer [SEP]"
tokenized_text = enc.tokenize(text)
masked_index = 8
tokenized_text[masked_index] = '[MASK]'
indexed_tokens = enc.convert_tokens_to_ids(tokenized_text)
segments_ids = [0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 1, 1]
tokens_tensor = torch.tensor([indexed_tokens])
segments_tensors = torch.tensor([segments_ids])
dummy_input = [tokens_tensor, segments_tensors]
config = AlbertConfig(torchscript=True)
model = AlbertModel(config)
model.eval()
model = AlbertModel.from_pretrained("albert-xxlarge-v2", torchscript=True)
traced_model = torch.jit.trace(model, [tokens_tensor, segments_tensors])
torch.jit.save(traced_model, "albert-xxlarge-v2.pt")

I successfully exported the model, but the second code cell threw out a RuntimeError:
Step 2

Code:

loaded_model = torch.jit.load("albert-xxlarge-v2.pt")
loaded_model.eval()

all_encoder_layers, pooled_output = loaded_model(dummy_input)

Outputs:

---------------------------------------------------------------------------

RuntimeError                              Traceback (most recent call last)

<ipython-input-7-2e61b1def246> in <module>()
      2 loaded_model.eval()
      3 
----> 4 all_encoder_layers, pooled_output = loaded_model(dummy_input)

/usr/local/lib/python3.6/dist-packages/torch/nn/modules/module.py in _call_impl(self, *input, **kwargs)
    720             result = self._slow_forward(*input, **kwargs)
    721         else:
--> 722             result = self.forward(*input, **kwargs)
    723         for hook in itertools.chain(
    724                 _global_forward_hooks.values(),

RuntimeError: forward() Expected a value of type 'Tensor' for argument 'input_ids' but instead found type 'list'.
Position: 1
Value: [tensor([[    2,    72,    23,  2170, 27674,    13,    60,     3,     4, 27674,
            23,    21, 10956,  7911,     3]]), tensor([[0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 1, 1]])]
Declaration: forward(__torch__.transformers.modeling_albert.AlbertModel self, Tensor input_ids, Tensor attention_mask) -> ((Tensor, Tensor))
Cast error details: Unable to cast Python instance to C++ type (compile in debug mode for details)

In order to ensure this is’t caused by my edit of the code, I also ran all the code there to export BERT to TorchScript and get same error at the same step.

Software Versions:

Name: torch
Version: 1.6.0+cu101

Name: transformers
Version: 3.0.2

Am I doing something wrong when converting the model?
Also, is there a way to reinitialise dummy_input without loading the pretrained model in Step 2?

Thanks for any advice

2

This link might be helpful in case you need still it.