open3d.ml.torch.layers.KNNSearch failed when more than 20 million points

[xleEmesent]

Checklist

• [X] I have searched for similar issues.
• [X] 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!!!

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

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

Additional information

No response

[WMQ1999]

It also appears when I try to query neighbour of 100k points in a 200k points point cloud. However, when I run this program again, everything is ok. So I try to run it iteratively, this bug happends after several epochs!

--------------------------------------update--------------------------------- After I update pytorch from 1.8.0 to 1.8.2, this bug seems to be fixed~~~