model load checkpoint pytorch

1# Additional information
2EPOCH = 5
3PATH = "model.pt"
4LOSS = 0.4
5
6##this is how you save model checkpoint
7
8torch.save({
9            'epoch': EPOCH,
10            'model_state_dict': net.state_dict(),
11            'optimizer_state_dict': optimizer.state_dict(),
12            'loss': LOSS,
13            }, PATH)
14

1##this is how you load the model from the checkpoint
2
3model = Net()
4optimizer = optim.SGD(net.parameters(), lr=0.001, momentum=0.9)
5
6checkpoint = torch.load(PATH)
7model.load_state_dict(checkpoint['model_state_dict'])
8optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
9epoch = checkpoint['epoch']
10loss = checkpoint['loss']
11
12model.eval()
13# - or -
14model.train()
15