Getting runtime error.

[abhinavsp0730]

Hi, @Ki6an it's great work. But while executing below code

from fastT5 import export_and_get_onnx_model
from transformers import AutoTokenizer

model_name = 't5-small'
model = export_and_get_onnx_model(model_name)

tokenizer = AutoTokenizer.from_pretrained(model_name)
t_input = "translate English to French: The universe is a dark forest."
token = tokenizer(t_input, return_tensors='pt')

tokens = model.generate(input_ids=token['input_ids'],
               attention_mask=token['attention_mask'],
               num_beams=2)

output = tokenizer.decode(tokens.squeeze(), skip_special_tokens=True)
print(output)

I'm getting this error.

RuntimeError: output with shape [5, 12, 1, 2] doesn't match the broadcast shape [5, 12, 2, 2]

[Ki6an]

@abhinavsp0730 thanks, that's a known issue in transformers library while exporting t5 to onnx. the issue can be fixed by following the guide in this notebook.

or you can make the following changes in the modeling_t5.py file of your local transformers library and save it:

it's just a two-line simple fix.
I've already created PR for this issue here

[abhinavsp0730]

@Ki6an thanks for the help. Is this method is supported for GPU or CPU?

[abhinavsp0730]

@Ki6an thanks for the help. Is this method is supported for GPU or CPU?

I'm guessing it's not using onnxruntime-gpu package.So is fastT5 compatable with Gpu if we install
onnxruntime-gpu in place of onnxruntime?

[Ki6an]

@abhinavsp0730 currently fastT5 doesn't support GPU, will be working on this in the future.
onnxruntime-gpu doesn't support quantization & for the onnx model to perform better than the PyTorch models on GPU need to implement IObinding.