voxel grid CUDA bugfix

rusty1s · Aug 31, 2018 · d85bc4f · d85bc4f
1 parent 920cc93
commit d85bc4f
Show file tree

Hide file tree

Showing 4 changed files with 3 additions and 18 deletions.
diff --git a/cuda/cluster.cpp b/cuda/cluster.cpp
diff --git a/cuda/grid_kernel.cu b/cuda/grid_kernel.cu
@@ -18,7 +18,7 @@ __global__ void grid_kernel(int64_t *cluster,
     for (ptrdiff_t d = 0; d < pos.sizes[1]; d++) {
       scalar_t p = pos.data[i * pos.strides[0] + d * pos.strides[1]] - start[d];
       c += (int64_t)(p / size[d]) * k;
-      k += (int64_t)((end[d] - start[d]) / size[d]);
+      k *= (int64_t)((end[d] - start[d]) / size[d]) + 1;
     }
     cluster[i] = c;
   }

diff --git a/setup.py b/setup.py
@@ -15,7 +15,7 @@
         CUDAExtension('grid_cuda', ['cuda/grid.cpp', 'cuda/grid_kernel.cu']),
     ]
 
-__version__ = '1.1.4'
+__version__ = '1.1.5'
 url = 'https://github.com/rusty1s/pytorch_cluster'
 
 install_requires = []

diff --git a/torch_cluster/__init__.py b/torch_cluster/__init__.py
@@ -1,6 +1,6 @@
 from .graclus import graclus_cluster
 from .grid import grid_cluster
 
-__version__ = '1.1.4'
+__version__ = '1.1.5'
 
 __all__ = ['graclus_cluster', 'grid_cluster', '__version__']