I guess the meaning of ‘vectorise’ is to calculate Hessian matrix for each row in input.
If it is true, I had the same problem.
This is the issue 5 that I posted.
I don’t think it is possible to do this efficiently with Pytorch.
So I switched to Jax.
This is the code 4 compares performance of Jax and Pytorch for Hassian diagonal vector calculation task. This might be what you need.
But I’m still figuring out if my approach is correct.