Does KNNSearch not support CUDA?

[dfldylan]

Checklist

• I have searched for similar issues.
• I have tested with the latest development wheel.
• I have checked the release documentation and the latest documentation (for main branch).

My Question

open3d==0.18, code

import torch
from open3d.ml.torch.layers import KNNSearch

# 假设已经有了数据和查询点，这里使用随机数据作为示例
points = torch.rand(100, 3, dtype=torch.float32).cuda()  # 点云数据，假设已经在 GPU 上
queries = torch.rand(10, 3, dtype=torch.float32).cuda()   # 查询点，假设已经在 GPU 上

# 创建 KNNSearch 层实例
knn_search = KNNSearch()

# 执行 k-NN 搜索
k = 5
neighbors_index, neighbors_distance = knn_search(points, queries=queries, k=k)

# `neighbors_index` 和 `neighbors_distance` 分别是查询点的最近邻索引和距离

result

/root/miniconda3/envs/gaussian_splatting/bin/python /workspace/gaussian/gaussian-splatting/utils/density.py 
Traceback (most recent call last):
  File "/workspace/gaussian/gaussian-splatting/utils/density.py", line 42, in <module>
    neighbors_index, neighbors_distance = knn_search(points=points, queries=queries, k=k)
  File "/root/miniconda3/envs/gaussian_splatting/lib/python3.7/site-packages/torch/nn/modules/module.py", line 1130, in _call_impl
    return forward_call(*input, **kwargs)
  File "/root/miniconda3/envs/gaussian_splatting/lib/python3.7/site-packages/open3d/ml/torch/python/layers/neighbor_search.py", line 382, in forward
    index_dtype=self.index_dtype)
  File "/root/miniconda3/envs/gaussian_splatting/lib/python3.7/site-packages/open3d/ml/torch/python/ops.py", line 642, in knn_search
    return_distances=return_distances))
  File "/root/miniconda3/envs/gaussian_splatting/lib/python3.7/site-packages/torch/_ops.py", line 143, in __call__
    return self._op(*args, **kwargs or {})
RuntimeError: KnnSearch does not support CUDA

Looking forward your answer