But how would you like to index the image?
Currently it has a shape of [1, 512, 512].
If you call image[0], you will get the pixel value of the one and only channel.
However, inside the loop you are increasing index by one, so that the next call would be image[1], which does not exist.
Could you try to explain, how the indexing should work or maybe draw a sketch?