You’re right those loops make it very slow but i need to use it somehow because i will try a function which uses the information of neighbours for each element.
By the way i tried your indexing but i noticed that when i change the upper limit of width and height to a very high value, it still gave me a result. I am using 4 images (3x256x256) with a batch size of 4 to at least learn how to implement a custom loss function. Can you explain why it still gives a result. Thanks



 superunification:

loss = torch.square(truth[:,:0:255,0:255] - pred[:,:0:255,0:255]).mean()