added assert

torch_cluster/grid.py

@@ -4,6 +4,9 @@ from .utils.ffi import grid
 def grid_cluster(pos, size, start=None, end=None):
     pos = pos.unsqueeze(-1) if pos.dim() == 1 else pos
 
+    assert pos.size(1) == size.size(0), (
+        'Last dimension of position tensor must have same size as size tensor')
+
     start = pos.t().min(dim=1)[0] if start is None else start
     end = pos.t().max(dim=1)[0] if end is None else end
     pos, end = pos - start, end - start