Segmentation fault using torch_cluster.knn when skipping batch numbers

🐛 Bug

torch_geometric gives a segfault when performing knn on a graph on CPU where the batch tensor skips a batch number.

To Reproduce

Example to reproduce the behavior:

import torch
from torch_cluster import knn_graph

points = torch.rand(2, 2)
batch = torch.tensor([0, 1])  
eidx = knn_graph(points, 1, batch)   # batch tensor with only consecutive numbers works fine

points = torch.rand(2, 2).cuda()
batch = torch.tensor([0, 2]).cuda()  
eidx = knn_graph(points, 1, batch)   # batch tensor skipping a number but on gpu runs fine

points = torch.rand(2, 2)
batch = torch.tensor([0, 2])
eidx = knn_graph(points, 1, batch)   # batch tensor skipping a number but on cpu gives a segfault

Expected behavior

It should work 😃

Environment

• OS: Tested on Mac OS 11.4 and Ubuntu 20.04.2
• Python version: 3.9.6 and 3.9.5
• PyTorch version: 1.9.0
• CUDA/cuDNN version: knn_cuda works
• GCC version:
• Any other relevant information:

Additional context

This specifically comes up for me when working with heterogeneous graphs with e.g. node types A and B, where occasionally some graphs may not have any nodes of type B, in which case the batch tensor for type B nodes will skip that graph number.

[rusty1s]

Yeah, sorry for the delay. We postponed the release to combine it with the PyTorch 1.11 release (which should be released today). I will try to re-build all wheels at the weekend.

[rusty1s]

Thanks for reporting. I fixed it in torch-cluster, see here.