PyTorch second derivative is zero everywhere
I'm working on a custom loss function that requires me to compute the trace of the hessian matrix of a neural network with respect to its inputs.
In one dimension, this reduces to the second derivative of the network with respect to its input. Suppose my network is u and my input is x. The problem is completely one-dimensional, so the relationship I'm trying to model is u(x) where u has one input and one output, and x is one-dimensional. However, I'm running with batches, so x &quot;comes in&quot; as a column vector and I'm expecting a column vector as output.
If we label the samples in the batch as x_1, x_2, ..., x_n, I'm thus interested in the following vectors:
     
In PyTorch, I have tried the following:
u = model(x)
d = grad(u, x, grad_outputs=torch.ones_like(u), create_graph=True)[0]
dd = grad(d, x, grad_outputs=torch.ones_like(d), retain_graph=True, create_graph=False)[0]

This works well for u' but u'' comes out as being zero everywhere:

I'm a bit confused on the whole notion of needing a vector jacobian product here to begin with. Should I view the computation as u mapping from R^n to R^n, where n is the size of my batch?
Any help is appreciated!

Toy example
I have tried to devise a toy example for you:
x = torch.rand(3, requires_grad=True).reshape(-1, 1)
u = x ** 2.

# du/dx = 2 * x
dx = torch.autograd.grad(u, x, torch.ones_like(u), retain_graph=True, create_graph=True)[0]

# d^2 u/dx^2 = 2
ddx = torch.autograd.grad(dx, x, torch.ones_like(u))[0]

These are the outputs of x, dx, ddx:
tensor([[0.5217],
        [0.0400],
        [0.5509]], grad_fn=&lt;ReshapeAliasBackward0&gt;)
tensor([[1.0434],
        [0.0800],
        [1.1018]], grad_fn=&lt;MulBackward0&gt;)
tensor([[2.],
        [2.],
        [2.]])

So grad is computing the gradient of x element-wise.
Should I view the computation as u mapping from R^n to R^n, where n is the size of my batch?
My answer is no, otherwise dx should have been a Jacobian matrix instead of a vector.
This works well for u' but u'' comes out as being zero everywhere
Given the limited context, the only thing I can think of is that the second derivative of model(x) with respect to x is zero.
