knn_graph returns empty edge_index when other gpus are used

[YanivDorGalron]

While using torch_cluster knn_graph function I get an empty edge_index depending on the order of the gpus used.

gpus != 0 will work only after gpu 0 has been used.

torch.version '2.4.0+cu121'
torch_cluster.version '1.6.2'

import torch
from torch_cluster import knn_graph
work_list = [0, 1]
will_not_work_list = [1, 0]
lst = will_not_work_list
for i in lst:
    # Example input tensor
    x = torch.rand((1000, 3))
    device = torch.device(f'cuda:{i}' if torch.cuda.is_available() else 'cpu')
    x = x.to(device)
    edge_index = knn_graph(x, k=5)
    print(edge_index.device)

Would highly appreciate your response

Thanks in advance,
Yaniv

[rusty1s]

This should be hopefully fixed once we build a new version including #213.

[rusty1s]

Can you install from source to see if this fixes your issues?