You can take the image display code (plt. calls before the output image) and move it to after optimizer.step() (also clamp, but don’t overwrite the input_img (so use detach().clamp()).
Best regards
Thomas