Pytorch Save & load the model 2

import torch
import torchvision.models as models

Saving and Loading Model Weights

model = models.vgg16(petrained=True)
torch.save(model.state_dict(),"model_weights.pth")

To load model weights ,you need to create an instance of the same model first,and then load the parameters using load_state_dict() method

model = models.vgg16() # we do not specify pretrained=True, i.e. do not load default weights
model.load_state_dict(torch.load('model_weights.pth'))
model.eval()

torch.save( model,'model.pth')

model = torch.load("model.pth")