Hi
I’m trying to calculate a gradient of a function that is  A*A*X*W0*W1 w.r.t A  in PyTorch. I have calculated the gradient manually but I’m not getting the same answer as the pytorch’s auto grad function. Can anyone help me with what I did wrong?
The manually calculated function is
def torch_grad(A, W0, W1, X):
  # Forward
  W0W1 = torch.mm(W0, W1)
  AX = torch.mm(A, X)
  AXW0W1 = torch.mm(AX, W0W1)
  AAXW0W1 = torch.mm(A, AXW0W1)

  # Backward / Gradients
  rows, cols = AAXW0W1.size()
  grad_AX = torch.mm(torch.eye(rows, cols), W0W1.t())
  grad_AA = torch.mm(grad_AX, X.t())
  grad_A = torch.mm(grad_AA, (2*A).t())
  return grad_A

Thanks!