It depends on what you want to do with the model later.
**For the purpose of resume training**
If you want to save a model to train it later, you will need more than the model's `state_dict` itself. You will also need to save the state of the optimizer, epochs, scores, etc.
state = {
'epoch': epoch,
'state_dict': model.state_dict(),
'optimizer': optimizer.state_dict(),
...
}
torch.save(state, filepath)
# load
model.load_state_dict(state['state_dict'])
optimizer.load_state_dict(state['optimizer'])
**For the purpose of inference**
Only the model's `state_dict` will suffice. But make sure you call the `eval` mode after you load the model so that batchnorm or dropout layers will work in eval mode instead of training mode.
torch.save(model.state_dict(), filepath)
# load
model.load_state_dict(torch.load(filepath))
model.eval()
Refer to the official [best practices guide](https://github.com/pytorch/pytorch/blob/761d6799beb3afa03657a71776412a2171ee7533/docs/source/notes/serialization.rst) if you need it.
It depends on what you want to do.
**Case # 1: Save the model to use it yourself for inference**: You save the model, you restore it, and then you change the model to evaluation mode. This is done because you usually have `BatchNorm` and `Dropout` layers that by default are in train mode on construction:
torch.save(model.state_dict(), filepath)
#Later to restore:
model.load_state_dict(torch.load(filepath))
model.eval()
**Case # 2: Save model to resume training later**: If you need to keep training the model that you are about to save, you need to save more than just the model. You also need to save the state of the optimizer, epochs, score, etc. You would do it like this:
state = {
'epoch': epoch,
'state_dict': model.state_dict(),
'optimizer': optimizer.state_dict(),
...
}
torch.save(state, filepath)
To resume training you would do things like: `state = torch.load(filepath)`, and then, to restore the state of each individual object, something like this:
model.load_state_dict(state['state_dict'])
optimizer.load_state_dict(state['optimizer'])
Since you are resuming training, **DO NOT** call `model.eval()` once you restore the states when loading.
**Case # 3: Model to be used by someone else with no access to your code**:
In Tensorflow you can create a `.pb` file that defines both the architecture and the weights of the model. This is very handy, specially when using `Tensorflow serve`. The equivalent way to do this in Pytorch would be:
torch.save(model, filepath)
# Then later:
model = torch.load(filepath)
This way is still not bullet proof and since pytorch is still undergoing a lot of changes, I wouldn't recommend it.