Hi,
I used below code to compute hessian matrix for a specific class. It seems to work but I need to double check. In below example I calculated the hessian matrix of output prediction for class 2 with respect to input tensor. I got an 4x4 hessian matrix in this way.
Can you please comment?
import torch

def jacobian(y, x, create_graph=False):
    jac = []
    flat_y = y.reshape(-1)
    grad_y = torch.zeros_like(flat_y)
    for i in range(len(flat_y)):
        grad_y[i] = 1.
        grad_x, = torch.autograd.grad(flat_y, x, grad_y, retain_graph=True, create_graph=create_graph)
        jac.append(grad_x.reshape(x.shape))
        grad_y[i] = 0.
    return torch.stack(jac).reshape(y.shape + x.shape)

def jacobian2(y, x, t, create_graph=False):
    flat_y = y.reshape(-1)
    grad_y = torch.zeros_like(flat_y)
    grad_y[t] = 1.
    grad_x, = torch.autograd.grad(flat_y, x, grad_y, retain_graph=True, create_graph=create_graph,allow_unused=True)
    return grad_x


def hessian(y, x, t):
    return jacobian(jacobian2(y, x, t, create_graph=True), x)


def f(x):
    return x * x * torch.arange(4, dtype=torch.float)


x = torch.ones(4, requires_grad=True) * 2

print(jacobian2(f(x), x, 2))
a = hessian(f(x), x, 2)
print(a)```