For a loss I had to compute the distance between each pair of points in a point cloud. I did it using the following method in a first time (x is a tensor of shape (Batch, nb_points, 3)):
dist = (x.unsqueeze(1) - x.unsqueeze(2)).pow(2).sum(-1).sqrt()
Giving a tensor dist of shape (nb_points, nb_points, 3). Using this implementation gave a float value in the loss but NaNs after back-propagation.
I replaced it by :
dist = torch.cdist(x,x)
and it worked properly.
My question is, why is it working in one case and not the other, is there anything different between my implementation and cdist? Is it related to gradients computation?
I expected the two implementation to have a similar behavior.