During the model training, I saved the parameters of the model in the following way:
torch.save(net.state_dict(),path)

After the training, I wanted to get a complete model, so I ran the following Python script:
import torch
import torch.nn as nn
import torchvision.models as models


class MobileNet(nn.Module):
    def __init__(self, num_classes, mode, pretrained=True):
        super(MobileNet, self).__init__()
        assert mode in ["mobilenet_v2", "mobilenet_v3_small", "mobilenet_v3_large"]
        self.num_classes = num_classes
        self.mode = mode

        if mode == "mobilenet_v2":
            model = models.mobilenet_v2(pretrained=pretrained)
            model.classifier[1] = nn.Linear(in_features=1280, out_features=self.num_classes, bias=True)

        elif mode == "mobilenet_v3_small":
            model = models.mobilenet_v3_small(pretrained=pretrained)
            model.classifier[3] = nn.Linear(in_features=1024, out_features=self.num_classes, bias=True)

        elif mode == "mobilenet_v3_large":
            model = models.mobilenet_v3_large(pretrained=pretrained)
            model.classifier[3] = nn.Linear(in_features=1280, out_features=self.num_classes, bias=True)

        self.model = model

    def forward(self, x):
        x = self.model(x)
        return x


if __name__ == "__main__":
    our_image = torch.rand(size=(2, 3, 224, 224))
    our_model = MobileNet(num_classes=4, mode="mobilenet_v2")
    our_model = torch.nn.DataParallel(our_model)
    our_model.load_state_dict(torch.load(path))
    torch.save(our_model, "./our_model.pt")


However, when I used the complete model I had saved, I got it wrong. The Python script used and the error is as follows:

1639019590(1)967×500 47.6 KB

I was rather puzzled by the error. There is an error indicating that I need to provide the script for building the model, but I have saved the complete model, so I don’t think I need the script for building the model.