They seem to be doing this https://github.com/f-dangel/backpack 9 and they too https://github.com/ChenAo-Phys/pytorch-Jacobian 9.
But they don’t support my use case directly. Any method based on this suitable here ?