I’m not sure why the above code cannot get the answer.
But I figure out another way:
x0 = torch.tensor([1., 2., 3.], requires_grad=True)
x1 = torch.tensor([4., 5., 6.], requires_grad=False)
x = torch.stack([x0, x1], dim=0)
y = x
d = torch.autograd.grad(y[0, 0], x0,  retain_graph=True, create_graph=True)
print(d)

output:
(tensor([1., 0., 0.]),)

But in this way, I have to take out the part of input x I need to calculate derivatives and then stack it with other parts of the input.
I used this method to calculate one column of the Jacobian matrix in a differentiable manner, but still out of memory. Any ideas for saving memory? The output size is (170, 2) and one part of inputs is (1, 768).