Hi,
This is because a[0] is a different Tensor from a. And that Tensor (that you just created) has not been used to compute the output.
You will have to do:
grad_a, = torch.autograd.grad(output, a)
grad_a_0 = grad_a[0]