Duplicate of Jacobian gradient matrix
The issue here is that flat_y[i] is a single number.
You should either give flat_y and keep doing what you do with the grad outputs so that they are both of the size of flat_y.
Or you can remove the grad outputs and use flat_y[i] so that they are both of size 1.