Hi,
I think the misconception here is between vector jacobian product vs jacobian vector product.
Reverse mode AD that all DL framework use for vector jacobian products. Where v is of the size of the input.
If you want to do jacobian vector product, (and so v is the same size as the input), you need to use a different trick. If you use pytorch nightly, we added functions to do this here 7.
Otherwise, you can use something based on this old gist: https://gist.github.com/apaszke/c7257ac04cb8debb82221764f6d117ad 1