fix jit bug

rusty1s · Mar 23, 2020 · aff91e0 · aff91e0
1 parent 6a8502c
commit aff91e0
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/torch_cluster/radius.py b/torch_cluster/radius.py
@@ -7,7 +7,7 @@
 @torch.jit.script
 def sample(col: torch.Tensor, count: int) -> torch.Tensor:
     if col.size(0) > count:
-        col = col[torch.randperm(col.size(0))][:count]
+        col = col[torch.randperm(col.size(0), dtype=torch.long)][:count]
     return col