Hey,
Note that you should not do x = torch.rand(batch_size, 3).requires_grad_(True).to("cuda:0") as the x you get here is not a leaf (it won’t have its .grad field populated when you call .backward()). You should do x = torch.rand(batch_size, 3, requires_grad=True, device="cuda:0")

I guess it is not possible to calculate the gradient of just part of a tensor. Am I right?

No you cannot ask for gradient from x[i]. Because x[i] is a new Tensor, different from x and that Tensor was not used to compute your function.