bugfix

torch_cluster/functions/grid.py

@@ -43,6 +43,7 @@ def grid_cluster(position, size, batch=None):
     func = get_func('grid', position)
     func(C, cluster, position, size, c_max)
     cluster = cluster.squeeze(dim=-1)
-    cluster = consecutive(cluster)
+    if not cluster.is_cuda:
+        cluster = consecutive(cluster)
 
     return cluster

torch_cluster/kernel/generic/kernel.cu

@@ -3,6 +3,7 @@
 #else
 
 void cluster_(grid)(THCState *state, int C, THCudaLongTensor *output, THCTensor *position, THCTensor *size, THCudaLongTensor *count) {
+  printf("drin");
 }
 
 #endif

torch_cluster/kernel/kernel.cu

@@ -2,7 +2,7 @@
 
 #include "kernel.h"
 
-#define cluster_(NAME) TH_CONCAT_4(cluster, NAME, _kernel_, Real)
+#define cluster_(NAME) TH_CONCAT_4(cluster_, NAME, _kernel_, Real)
 
 #include "generic/kernel.cu"
 #include "THCGenerateAllTypes.h"

torch_cluster/src/generic/cuda.c

@@ -3,7 +3,7 @@
 #else
 
 void cluster_(grid)(int C, THCudaLongTensor *output, THCTensor *position, THCTensor *size, THCudaLongTensor *count) {
-  /* return cluster_kernel_(grid)(state, C, output, position, size, count); */
+  cluster_kernel_(grid)(state, C, output, position, size, count);
 }
 
 #endif