KNN selects fewer number of neighbours

[dadin4vcc]

🐛 Describe the bug

Hi, Thanks for the great work with PyG! As reported in this issue, the knn_graph function gives too few neighbours (possiby only on GPU). It seems the same thing goes for the knn-function The proposed fix however works fine.

Environment

• PyG version: 2.2.0
• PyTorch version: 1.12.1
• OS: Linux
• Python version: 3.9.16
• CUDA/cuDNN version: 10.2.89
• How you installed PyTorch and PyG (conda, pip, source): conda
• Any other relevant information (e.g., version of torch-scatter):

[rusty1s]

The fix in https://github.com/rusty1s/pytorch_cluster/commit/113e8a6abbe3c2856b1656f111bb4716770ddfdc should work for both knn and knn_graph. Do you have a small example to reproduce?

[JeppeLS]

I'm still experiencing this bug on these versions

torch                     2.2.2
torch_cluster             1.6.3+pt22cu121
torch_geometric           2.5.2
torch_sparse              0.6.18+pt22cu121

A small example to reproduce

import torch
from torch_geometric.nn import knn

x = torch.tensor([[0, 0], [100000, 100000]]).cuda().double()
y = torch.tensor([[1, 1]]).cuda().double()

print(knn(x, y, 2))

x = x.cpu()
y = y.cpu()

print(knn(x, y, 2))

yields

tensor([[0],
        [0]], device='cuda:0')
tensor([[0, 0],
        [0, 1]])