Fix bug

Author: nicolas-chaulet

torch_points_kernels/__init__.py

@@ -1,4 +1,13 @@
 from .torchpoints import *
 from .knn import knn
+from .cluster import region_grow
 
-__all__ = ["ball_query", "furthest_point_sample", "grouping_operation", "three_interpolate", "three_nn", "knn"]
+__all__ = [
+    "ball_query",
+    "furthest_point_sample",
+    "grouping_operation",
+    "three_interpolate",
+    "three_nn",
+    "knn",
+    "region_grow",
+]

torch_points_kernels/cluster.py

@@ -97,7 +97,7 @@ def region_grow(
         if len(label_clusters):
             remaped_clusters = []
             for cluster in label_clusters:
-                cluster = cluster.to(pos.device)
+                cluster = torch.tensor(cluster).to(pos.device)
                 remaped_clusters.append(local_ind[cluster])
             clusters[l.item()] = remaped_clusters
 

torch_points_kernels/torchpoints.py

@@ -6,7 +6,6 @@
 
 import torch_points_kernels.points_cpu as tpcpu
 from .knn import knn
-from .cluster import region_grow
 
 if torch.cuda.is_available():
     import torch_points_kernels.points_cuda as tpcuda