082a83b8eb9177051864:

B[0, :, :] = B[0, :, :].mm(A)


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)
loss = z.sum()
loss.backward()

When I assign updated B to z, I am not mutating B so autograd has access to original B.
Best regards
Nik