Solved by mariosasko in post #2 
              
              
                Replace the line self.R = torch.tensor([[0.01]], requires_grad=True) with self.R = torch.nn.Parameter(torch.tensor([[0.01]])). After that modification, the parameters will be properly registered, so they will reside in the model’s state dict. And you should remove your own implementation of the meth…