[FIXED] Convert a pth pytorch file to an onnx model

Issue

I'm trying to convert a PyTorch model(pth file containing weights) to an onnx file then to a TensorFlow model since I work on TensorFlow. to then fine-tune it. This is my attempt so far. I keep however getting errors.enter image description here I think the problem is that the weights are for a vision transformer. But I haven't figure out what type of model to use to convert it. I'm assuming a CRNN but if there is an easier way I would love to know. PS: I did load the pth file to my drive. the path is correct

from torch.autograd import Variable

import torch.onnx
import torchvision
import torch
import onnx
import torch.nn as nn

dummy_input = torch.randn(1, 3, 224, 224)
file_path='/content/drive/MyDrive/VitSTR/vitstr_base_patch16_224_aug.pth'

model = torchvision.models.vgg16()

model.load_state_dict(torch.load(file_path))

model.eval()

torch.onnx.export(model, dummy_input, "vitstr.onnx")

Solution

Thank you all. I used the same architecture as the one in the model and it worked.

Answered By - Roua Rouatbi