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
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