How to save a model from a previous epoch?

Recommended approach

path = os.path.join(SAVE_DIR, 'model.pth')
torch.save(MODEL.cpu().state_dict(), path) # saving model
MODEL.cuda() # moving model to GPU for further training
1 Like