Convert a CUSTOM image segmentation PyTorch model into ONNX

I followed this tutorial to convert a pretrained U-NET: https://developer.nvidia.com/blog/speeding-up-deep-learning-inference-using-tensorrt/
However, I could not convert a custom model. I have the .pt file and the Model class, and I am using torch.load(‘unet.pt’) instead of torch.hub.load (which loads the model from github repo).
I am getting this error:
AttributeError: ‘collections.OrderedDict’ object has no attribute 'training’

Anyone was able to solve it or convert a custom Unet? Appreciated.

Hi,

Please check this GitHub to convert the .pth file into TensorRT.

Thanks.