100% codecov

test/dense_grid.json

@@ -37,6 +37,16 @@
     "batch": [0, 0, 0, 0, 0, 1, 1],
     "expected": [0, 5, 1, 0, 2, 6, 9],
     "expected_C": 6
+  },
+  {
+    "name": "Batch with start/end parameter",
+    "position": [[0, 0], [11, 9], [2, 8], [2, 2], [8, 3], [1, 1], [6, 6]],
+    "size": [5, 5],
+    "batch": [0, 0, 0, 0, 0, 1, 1],
+    "start": 0,
+    "end": 20,
+    "expected": [0, 9, 1, 0, 4, 16, 21],
+    "expected_C": 16
   }
 ]
 

test/sparse_grid.json

@@ -31,5 +31,11 @@
     "batch": [0, 0, 0, 0, 0, 1, 1],
     "expected": [0, 3, 1, 0, 2, 4, 5],
     "expected_batch": [0, 0, 0, 0, 1, 1]
+  },
+  {
+    "name": "Position tensor",
+    "position": [[[0, 0], [9, 9]], [[0, 0], [9, 0]]],
+    "size": [5, 5],
+    "expected": [[0, 2], [0, 1]]
   }
 ]

torch_cluster/functions/grid.py

@@ -48,12 +48,13 @@ def _fixed_cluster_size(position, size, batch=None, end=None):
     if end is None:
         return _minimal_cluster_size(position, size)
 
-    eps = 0.000001  # Model [start, end).
+    eps = 0.000001  # Simulate [start, end) interval.
     if batch is None:
         cluster_size = ((end / size).float() - eps).long() + 1
     else:
         cluster_size = ((end / size[1:]).float() - eps).long() + 1
-        cluster_size = torch.cat([batch.max() + 1, cluster_size], dim=0)
+        max_batch = cluster_size.new(1).fill_(batch.max() + 1)
+        cluster_size = torch.cat([max_batch, cluster_size], dim=0)
 
     return cluster_size