Solved by albanD in post #2 
              
              
                To add learnable parameters to an nn.Module, you should use nn.Parameter type: 
self.W = nn.Parameter(torch.randn(dim, 1).type(dtype)) 
Then you need to make sure that this is passed to your optimizer with the other parameters of your network: 
optimizer = optim.SGD(model.parameters() + criterion.pa…