sorry i am not getting the “TheModelClass”
model = TheModelClass(*args, **kwargs)
do you mean this in my .pth file
class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.conv1 = nn.Conv2d(3, 32, 3)
        self.pool = nn.MaxPool2d(2,    2)
        self.conv2 = nn.Conv2d(32, 64, 3)
        self.fc1 = nn.Linear(64*14*14,30)
        self.fc2 = nn.Linear(30, 3)
        #self.softmax = nn.Softmax(dim=1)

    def forward(self, x):
        x = self.pool(F.relu(self.conv1(x)))  # output size: [batch_size, 32, 255, 255]
        x = self.pool(F.relu(self.conv2(x)))  # output size: [batch_size, 64, 126, 126]
        
        x = x.view(-1,64*14*14)  # output size: [batch_size, 64*126*126]
        x = F.relu(self.fc1(x))
        x = self.fc2(x)
        #x=self.softmax(x)
        return x


x = torch.randn(1, 3, 64, 64) #(batch size or #of images,channels RGB,width,height)
model = Net()
output = model(x)