Mismatch error ,, jacobian gradient matrix

i have input image with size(1,1,28,28) and output image y with the same size. y= cnn(x)
and i want to calculate the gradient between every output pixel image with respect to every pixel from input image but i have : RuntimeError: Mismatch in shape: grad_output[0] has a shape of torch.Size([784]) and output[0] has a shape of torch.Size([]).
my cod is :
def jacobian(y, x, create_graph=False):

jac = []
y = cnn(x)

flat_y = y.reshape(-1) 
flat_x = x.reshape(-1)
#y= cnn(flat_x)
grad_y = torch.zeros_like(flat_y) 

for i in range(len(flat_y)): 
    for j in range(len(flat_x)):
        grad_y[i] = 1. 
        grad_x,= torch.autograd.grad(flat_y[i] ,flat_x[j],grad_y,retain_graph=True,create_graph= create_graph,allow_unused=True)

        grad_y[i] = 0.

return torch.stack(jac).reshape(y.shape + x.shape)  

print( jacobian (y,x,create_graph = True) )

could someone help me why i have this error?

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.