Save and load checkpoint pytorch

Question

i make a model and save the configuration as:

def checkpoint(state, ep, filename='./Risultati/checkpoint.pth'):  
    if ep == (n_epoch-1):
        print('Saving state...')
        torch.save(state,filename)
checkpoint({'state_dict':rnn.state_dict()},ep) 

and then i want load this configuration :

state_dict= torch.load('./Risultati/checkpoint.pth')
    rnn.state_dict(state_dict)

when i try, this is the error:

Traceback (most recent call last):
File "train.py", line 288, in <module>
rnn.state_dict(state_dict)
File "/home/marco/.local/lib/python3.6/site-packages/torch/nn/modules/module.py", line 593, in state_dict
destination._metadata[prefix[:-1]] = dict(version=self._version)
AttributeError: 'dict' object has no attribute '_metadata'

where i do wrong?

thx in advance

Answer 1

You need to load rnn.state_dict() stored in the dictionary you loaded:

rnn.load_state_dict(state_dict['state_dict'])

Look at load_state_dict method for more information.