blackberry:

Would be great if PyTorch provided and option to torch.save/torch.load to save load either the whole model or the state_dict.


The recommended way of storing the state_dict should be working fine or are you seeing any issues with it? As explained in the Serialization semantics 23, storing and loading the model directly could be breaking in various ways due to the way Python pickles the model.