I just tried to register the buffer as a torch.Tensor:
init_mem_bias = torch.Tensor(self.N, self.M)
self.register_buffer(‘mem_bias’, init_mem_bias)
This triggered invalid types later in the program where a Variable is expected.
File “C:\Users\frede\Documents\GitHub\Homework3\ntm\memory.py”, line 177, in _similarity
w = F.softmax(β * F.cosine_similarity(self.memory + 1e-16, k + 1e-16, dim=-1), dim=1)
File “C:\Users\frede\Anaconda3\lib\site-packages\torch\nn\functional.py”, line 1639, in cosine_similarity
w12 = torch.sum(x1 * x2, dim)
File “C:\Users\frede\Anaconda3\lib\site-packages\torch\tensor.py”, line 321, in mul
return self.mul(other)
TypeError: mul received an invalid combination of arguments - got (Variable), but expected one of:

(float value)
didn’t match because some of the arguments have invalid types: (!Variable!)
(torch.FloatTensor other)
didn’t match because some of the arguments have invalid types: (!Variable!)

At this stage, since I have a report to submit, I will just have to do without saving the model. It might be the complicated structure of the Neural Turing Machine, because I normally don’t have any  issues with save and load state_dict.
Thanks a lot for your input anyway.