better dense return type

Author: rusty1s

setup.py

@@ -2,7 +2,7 @@
 
 from setuptools import setup, find_packages
 
-__version__ = '0.2.2'
+__version__ = '0.2.3'
 url = 'https://github.com/rusty1s/pytorch_cluster'
 
 install_requires = ['cffi', 'torch-unique']

test/dense_grid.json

@@ -36,7 +36,7 @@
     "size": [5, 5],
     "batch": [0, 0, 0, 0, 0, 1, 1],
     "expected": [0, 5, 1, 0, 2, 6, 9],
-    "expected_C": 6
+    "expected_C": 12
   },
   {
     "name": "Batch with start/end parameter",
@@ -46,7 +46,7 @@
     "start": 0,
     "end": 20,
     "expected": [0, 9, 1, 0, 4, 16, 21],
-    "expected_C": 16
+    "expected_C": 32
   }
 ]
 

torch_cluster/__init__.py

@@ -1,5 +1,5 @@
 from .functions.grid import sparse_grid_cluster, dense_grid_cluster
 
-__version__ = '0.2.2'
+__version__ = '0.2.3'
 
 __all__ = ['sparse_grid_cluster', 'dense_grid_cluster', '__version__']

torch_cluster/functions/grid.py

@@ -88,9 +88,4 @@ def dense_grid_cluster(position, size, batch=None, start=None, end=None):
     position, size = _preprocess(position, size, batch, start)
     cluster_size = _fixed_cluster_size(position, size, batch, end)
     cluster, C = _grid_cluster(position, size, cluster_size)
-
-    if batch is None:
-        return cluster, C
-    else:
-        C = C // cluster_size[0]
-        return cluster, C
+    return cluster, C