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.