Thank you for your quick response!
I think A.backward(gradient=torch.ones_like(A))  is the same as A.sum().backward(). Consequently, you don’t compute the derivative of every element of A w.r.t. every element of B, but the derivative of the batch-wise sum of B w.r.t. A. Consequently, you get a gradient of shape (4, 10), not (4, 20, 10).