0

I am converting pytorch model into Tensorflow. For that Pytorch-ONNX-Tensorflow this are the steps I am following, but while exporting Pytorch model into ONNX getting the following error: ModuleNotFoundError: No module named 'models'

Below is the code snippet:

trained_model = Net()
trained_model.load_state_dict(torch.load('best.pt'))
dummy_input = Variable(torch.randn(1, 1, 28, 28))
torch.onnx.export(trained_model, dummy_input, "best.onnx")

Any help would be appreciated. Thanks!

Shivani Akotkar
  • 1
  • 2
  • 1
    Could you paste the rest of the code. Because the error probably coming from another part of the code. – yakhyo Dec 03 '21 at 05:01
  • Also post the full error traceback – DerekG Dec 03 '21 at 13:59
  • model_path = '/content/yolov5s.pt' # Load the weights state_dict = torch.load(model_path) – Shivani Akotkar Dec 05 '21 at 06:19
  • @DerekG : 2 frames /usr/local/lib/python3.7/dist-packages/torch/serialization.py in find_class(self, mod_name, name) 873 def find_class(self, mod_name, name): 874 mod_name = load_module_mapping.get(mod_name, mod_name) 875 return super().find_class(mod_name, name) 876 877 # Load the data (which may in turn use `persistent_load` to load tensors) ModuleNotFoundError: No module named 'models' – Shivani Akotkar Dec 05 '21 at 06:23

0 Answers0