Saving and loading SGD optimizer

I’m trying to save SGD Optimizer state using, 'optim.pth') and then loading it with optimizer.load_state_dict(torch.load('optim.pth'))
But after loading the state, optimizer.step() gives error for in for following line

if momentum != 0: param_state = self.state[p]
This is probably because it doesn't save the dictionary used by [optimizer]( So I tried using `, 'optim.pth')` but then the network simply doesn't train at all. So what' the correct way of saving and loading SGD optimizer (while using momentum) ?

Just found out that this is because it doesn’t save the dictionary used by [optimizer] ( as it is while using optimizer.state_dict(). Before saving the optimizer, it’s state variable is defaultdict(<type 'dict'>, {}) but when it is loaded back it changes to a simple dict {}

So for now I’m training the network on a random batch before saving the state so that it doesn’t have to store a blank dict.

Could you, please explain then how to save and reload an optimizer.

I understand that if one is saving the model, the optimizer must also be saved otherwise the learning_rate value inside the optimizer is lost.


A quick workaround that fixed it for me was doing optimizer.state = defaultdict(dict, optimizer.state) after doing optimizer.load_state_dict(torch.load('optim.pth')).