Awesome, thank you!
The expand is not necessary (if I don’t miss anything). And your idea extended to a batch of coordinates looks like this:
# y shape: BS x N_COORDS x 2
# dist_gt shape: BS x N_COORDS x N_COORDS x 2
dist_gt = torch.zeros(bs, n_lm, n_lm, 2)
dist_gt[:,:,:,0] = y[:,:,0].view(bs,1,-1) - y[:,:,0].view(bs,-1,1)
dist_gt[:,:,:,1] = y[:,:,1].view(bs,1,-1) - y[:,:,1].view(bs,-1,1)

It is much faster than looping over the tensor manually (~5s for the whole dataset compared to 12 min)