img_output = model.forward(img_input)

This is Incorrect. Has to be
img_output = model(img_input)

You are treating this as a scalar gradient, and hence you treat grad_img_input as getting wrt a scalar variable 1.