Bug/regression in radius function

[leachim]

This is related to #217 , but seems not to be restricted to ARM64. I am having the same issue on a x86_64 cluster environment.

Traceback (most recent call last):
  File "/cluster/home/michaes/hyperion/debug.py", line 13, in <module>
    assign_index = radius(x, y, 1.5, batch_x, batch_y)
                   ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/cluster/home/michaes/.miniforge/envs/hyperion/lib/python3.11/site-packages/torch_cluster/radius.py", line 82, in radius
    return torch.ops.torch_cluster.radius(x, y, ptr_x, ptr_y, r,
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/cluster/home/michaes/.miniforge/envs/hyperion/lib/python3.11/site-packages/torch/_ops.py", line 854, in __call__
    return self_._op(*args, **(kwargs or {}))
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
RuntimeError: x.dim() == 2 INTERNAL ASSERT FAILED at "csrc/cpu/radius_cpu.cpp":13, please report a bug to PyTorch. Input mismatch

Can be reproduced with this minimal example:

from torch_cluster import radius
import torch

x = torch.tensor([[-1.0, -1.0], [-1.0, 1.0], [1.0, -1.0], [1.0, 1.0]])
batch_x = torch.tensor([0, 0, 0, 0])
y = torch.tensor([[-1.0, 0.0], [1.0, 0.0]])
batch_y = torch.tensor([0, 0])
assign_index = radius(x, y, 1.5, batch_x, batch_y)

[leachim]

appears to be fixed in the very latest releases, so closing this.