Continuing the discussion from Pytorch most efficient Jacobian / Hessian calculation:
This is the fastest way I can tell
suppose you have a loss function of x
        loss = f(x)
        first_drv = torch.zeros(batch_size, x_dim)
        hessian = torch.zeros(batch_size, x_dim, x_dim)
        for n in range(batch_size):
            first_drv[n] = torch.autograd.grad(loss[n], dz,
                                                     create_graph=True, retain_graph=True)[0][n]
            for i in range(x_dim):
                hessian[n][i] = torch.autograd.grad(first_drv[n][i], dz,
                                                        create_graph=True, retain_graph=True)[0][n]