You could try an earthmovers distance (https://en.m.wikipedia.org/wiki/Earth_mover%27s_distance) in 2d or 3d over the image? For example you could do this, but call sequentially (https://discuss.pytorch.org/t/implementation-of-squared-earth-movers-distance-loss-function-for-ordinal-scale/107927/2) The idea would be something like this (untested and written on my cell phone): def cumsum_3d(a): a = torch.cumsum(a, -1) a = torch.cumsum(a, -2) a = torch.cumsum(a, -3) return a def norm_3d(a): return a / torch.sum(a, dim=(-1,-2,-3), keepdim=True) def emd_3d(a, b): a = norm_3d(a) b = norm_3d(b) return torch.mean(torch.square(cumsum_3d(a) - cumsum_3d(b)), dim=(-1,-2,-3)) This should also work with batched data. I would also try normalizing the images first (so they each sum to 1) unless you want to account for changes in intensity.