Jacobian gradient matrix

Hi ,
i have this code to calculate the image gradient output w.r.t input (for every pixel from output image w.r.t each pixel in input image) output image and input image have the same size (28 *28)
but i have missmatch error. Mismatch in shape: grad_output[0] has a shape of torch.Size([784]) and output[0] has a shape of torch.Size([]).

my code

def jacobian(y, x, create_graph=False):

image = Image.open('/Users/lionardo/Desktop/4_.png')

x = TF.to_tensor(image)
x = x.unsqueeze(0)
x.requires_grad_(True)

y = cnn(x)
jac = []
flat_y = y.reshape(-1)
flat_x = x.reshape(-1)
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 = True )

jac.append(grad_x.reshape(x.shape))  

grad_y[i] = 0. 
   
#return torch.stack(jac).reshape(y.shape + x.shape) 

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

# return jacobian(jacobian(y, x, create_graph=True))

print (jacobian)