Hi @sakuraiiiii,
TLDR: You need to use the create graph functionality in the grad function.
That is, changing from the below line



 sakuraiiiii:

grad_x = grad(output.sum(), x, retain_graph=True)[0]


to
grad_x = grad(output.sum(), x, create_graph=True)[0]

Explanation:
Create graph=True makes the graph of the derivative to be constructed, allowing you to compute higher order derivatives. The grad tensors hence will get a backward function attached to them.