Solved by Nikronic in post #2 
              
              
                Hi, 
In this line, you are updating B variable but autograd needs all previous values to compute the backpropagation. 
You can do something like this to ensure you are not losing the value of B: 
import torch

A = torch.nn.Parameter(torch.randn(3, 3))
B = torch.randn(2, 4, 3)
z = B[0, :, :].mm(A)
l…