bruce:

corrected_im = corrected_im - lr*corrected_im.grad


I think the problem is this line: It changes the variable corrected_im to the result of the substraction. But since this substraction is done in a no_grad block, the result does not require gradient.
I think you want to change it inplace no?
corrected_im -= lr * corrected_im.grad
Also I don’t think .zeros_() is a thing, you want .zero_() no?