I am trying to calculate the trace of Hessian of a sample neural network model.

I am using `torch.autograd.functional.hessian ` which basically returns a tuple of tuple of Tensors

for simplicity, let’s assume we have a linear model with 2 inputs and 2 outputs this will be my Hessian in this case for a random input and chosen loss function.

``````H =
((tensor([[[[ 0.1600,  0.1540],
[-0.1600, -0.1540]],

[[ 0.1540,  0.1482],
[-0.1540, -0.1482]]],

[[[-0.1600, -0.1540],
[ 0.1600,  0.1540]],

[[-0.1540, -0.1482],
tensor([[[ 0.1972, -0.1972],
[ 0.1898, -0.1898]],

[[-0.1972,  0.1972],
How can I calculate the trace of Hessian in this case? `torch.trace` didn’t work for me any ideas or tips?