Well your nested loops aren’t speeding it up. It’s slow because you’re doing 65,000+ pixel-wise MSEs, which I’m guessing is making your autograd graph extremely long. Not only that but there’s no parallelism possible for these calculations.
Why not just diff between truth and pred wholesale?
loss = torch.square(truth[:,:0:255,0:255] - pred[:,:0:255,0:255]).mean()
or if your images are already 256 pix wide and high (I’m assuming 255 is an off-by-one error?), then just
loss = torch.square(truth - pred).mean()