How to efficiently find correspondences between two point sets without nested for loop in Pytorch?
I now have two point sets (tensor) A and B that shape like

A.size() >>(50, 3) , example: [ [0, 0, 0], [0, 1, 2], ..., [1, 1, 1]]

B.size() >>(10, 3)  

where the first dimension stands for number of points and the second dim stands for coordinates (x,y,z)

To some extent, the question could also be simplified into " Finding common elements between two tensors ". Is there a quick way to do this without nested loop ?

You can quickly compute all the 50x10 distances using:

d2 = ((A[:, None, :] - B[None, ...])**2).sum(dim=2)


Once you have all the pair-wise distances, you can select "similar" ones if the distance does not exceed a threshold thr:

(d2 &lt; thr).nonzero()  


returns pairs of a-idx, b-idx of "similar" points.

If you want to match the points exactly, you can do instead:

((a[:, None, :] == b[None, ...]).all(dim=2)).nonzero()

