Computing sum of square of diagonal entries of Hessian

[atif-hassan]

Hi!!

First and foremost I would like to say that PyHessian is an incredible package and I am really thankful to the team for open-sourcing it!

I am a PhD student currently studying how the loss landscape behaves with change in neural network depth. For my study, I need to compute the sum of the square of the diagonal entries of the Hessian rather than just the sum (trace). Is there a way to do this with the package?