Incorrect square distance function

[mbgniuxya]

Hi, thanks for your amazing work.

I used your code on my customized dataset, which I down sampled using voxelization with voxel size of 0.05. In the function query_ball_points it gives incorrect indices and causes out of bound error in cuda. The reason is in the first set abstraction layer it tries to search the point's neighbours within region 0.05, but my data is already down sampled, which mean most of the time the interested point will have no neighbours. Even though, the shortest distance given from the distance matrix (N * M) should be 0, so the first element in the index would be that point itself. However, this is not the case. I think the error resides in the function square_distance, although I haven't identified which line of code is the problem, but the distance calculated by this function is incorrect.

[mbgniuxya]

suggested correction:

diff = src.unsqueeze(2).expand(B, N, M, C) - dst.unsqueeze(1).expand(B, N, M, C)
dist = torch.sum(diff ** 2, -1)

The original function gives negative value for square distance.