ONNX T5 - Decoding seq2seq tokens

I am trying to debug an issue decoding text I generate using a finetuned T5 model that I have created yet I continue to encounter the TypeError shown below.

Traceback (most recent call last):
  File "generate-mirage-onnx.py", line 258, in <module>
    output_text, output_logits = flan_t5(prompt, max_length=512, temperature=0.)
  File "/opt/conda/envs/accelerate/lib/python3.7/site-packages/torch/nn/modules/module.py", line 1194, in _call_impl
    return forward_call(*input, **kwargs)
  File "generate-mirage-onnx.py", line 112, in forward
    return self.tokenizer.decode(new_tokens), new_logits
  File "/opt/conda/envs/accelerate/lib/python3.7/site-packages/transformers/tokenization_utils_base.py", line 3480, in decode
    **kwargs,
  File "/opt/conda/envs/accelerate/lib/python3.7/site-packages/transformers/tokenization_utils_fast.py", line 549, in _decode
    text = self._tokenizer.decode(token_ids, skip_special_tokens=skip_special_tokens)
TypeError: argument 'ids': 'float' object cannot be interpreted as an integer

I am quite familiar with the tokenizer functionality of Transformers but it appear that there is something about seq2seq models that I am missing. Any tips would be greatly appreciate!

My full script can be found here.

@nielsr any tips or suggestions would be greatly appreciated.