Failed to convert torchscript model use torch.jit.script

jit
1

Hi Community,

I have came across a issue that throw me an error when convert model:
AssertionError: Unsupported annotation typing.Union[torch.FloatTensor, NoneType] could not be resolved because <class ‘torch.FloatTensor’> could not be resolved.

This looks to me that the script does not import torch package correctly, but I checked all the code I see torch is imported correctly

system ubuntu-18.04
torch 1.8.0
torchvision 0.9.0

Anyone come across similar issue?

Thanks