Thanks for your quick reply.
mem_bias in this case is just an initialization variable, so it is not meant to be a trainable Parameter.
I tried not to register any buffer, but I get the following error message when I try to load state_dict.
C:/Users/frede/Documents/GitHub/Homework3/Evaluate_NTM_LSTM.py
Traceback (most recent call last):
C:\Users\frede\Documents\GitHub\Homework3\copy\copy-task_NTM-LSTM-1000-batch-1000.model
File “C:/Users/frede/Documents/GitHub/Homework3/Evaluate_NTM_LSTM.py”, line 56, in 
loadNTM(model_NTM_LSTM, files_NTM_LSTM)
File “C:\Users\frede\Documents\GitHub\Homework3\train.py”, line 89, in loadNTM
model.net.load_state_dict(state_dict)
File “C:\Users\frede\Anaconda3\lib\site-packages\torch\nn\modules\module.py”, line 526, in load_state_dict
raise KeyError(‘missing keys in state_dict: “{}”’.format(missing))
KeyError: ‘missing keys in state_dict: “{‘memory.mem_bias’, ‘ntm.memory.mem_bias’, ‘ntm.heads.1.memory.mem_bias’, ‘ntm.heads.0.memory.mem_bias’}”’
Any other ideas about how to be able to save and load it in this setting?