AlphaBetaGamma96:

re-use the same graph but just for different inputs?


The whole memory usage form the graph comes from intermediary results. So re-using the graph wouldn’t help.
I’m afraid that extracting the diagonal of the Hessian cannot be done more efficiently in general.
You could reduce the memory usage by doing a custom function that drops the non-diagonal elements earlier but that won’t be a huge change.
You might want to check packages like this one https://backpack.pt/ 16 that provides approximate method to get the diagonal elements. If you are happy with an approximate value for your use case, it will be the fastest you can do!