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 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 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.