Convert a CUSTOM image segmentation PyTorch model into ONNX

I followed this tutorial to convert a pretrained U-NET:
However, I could not convert a custom model. I have the .pt file and the Model class, and I am using torch.load(‘’) 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.


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