open3d.ml.torch.layers.KNNSearch failed when more than 20 million points (int type error)

[xleEmesent]

Checklist

• [X] I have searched for similar issues.
• [X] For Python issues, I have tested with the latest development wheel.
• [X] I have checked the release documentation and the latest documentation (for master branch).

Describe the issue

Hi, all. I found out that open3d.ml.torch.layers.KNNSearch failed to generate the results when there are about more than 20 million points. The program just exits at ans = nsearch(points, queries, K) without warning message.

Is there any workaround or anything else I have missed here? Thanks in advance!!!

I am suspecting in the cuda supported cpp file at kine 244 the integer type has limited number of points can be queried. https://github.com/isl-org/Open3D/blob/746247447b70b24d1084a218ae4b07b373a95cce/cpp/open3d/core/nns/KnnSearchOps.cu#L244

int num_points = points.GetShape(0);
int num_queries = queries.GetShape(0);

Steps to reproduce the bug

import open3d.ml.torch as ml3d
import torch

points = torch.from_numpy(xyz)
queries = torch.from_numpy(xyz)
   
nsearch = ml3d.layers.KNNSearch(return_distances=True)
ans = nsearch(points, queries, K)
    
neighbors_index = ans.neighbors_index.reshape(xyz.shape[0],K)
neighbors_distance = ans.neighbors_distance.reshape(xyz.shape[0],K)

Error message

No response

Expected behavior

No response

Open3D, Python and System information

- Operating system: Ubuntu 20.04 / macOS 10.15 / Windows 10 64-bit
- Python version: Python 3.8 / output from `import sys; print(sys.version)`
- Open3D version: output from python: `print(open3d.__version__)`
- System architecture: x86 / arm64 / apple-silicon / jetson / rpi
- Is this a remote workstation?: yes or no
- How did you install Open3D?: pip / conda / build from source
- Compiler version (if built from source): gcc 7.5 / clang 7.0

Additional information

No response