bugfix batch == None

torch_cluster/fps.py

@@ -26,6 +26,9 @@ def fps(x, batch=None, ratio=0.5, random_start=True):
         >>> sample = fps(pos, batch)
     """
 
+    if batch is None:
+        batch = x.new_zeros(x.size(0), dtype=torch.long)
+
     assert x.is_cuda
     assert x.dim() <= 2 and batch.dim() == 1
     assert x.size(0) == batch.size(0)
@@ -33,9 +36,6 @@ def fps(x, batch=None, ratio=0.5, random_start=True):
 
     x = x.view(-1, 1) if x.dim() == 1 else x
 
-    if batch is None:
-        batch = x.new_zeros(x.size(0), dtype=torch.long)
-
     op = fps_cuda.fps if x.is_cuda else None
     out = op(x, batch, ratio, random_start)