fix jit bug

aff91e0e · rusty1s · 6a8502cb · aff91e0e
Commit aff91e0e authored Mar 23, 2020 by rusty1s
Hide whitespace changes
Inline Side-by-side

Showing with 1 addition and 1 deletion

torch_cluster/radius.py torch_cluster/radius.py +1 -1

No files found.
--- a/torch_cluster/radius.py
+++ b/torch_cluster/radius.py
@@ -7,7 +7,7 @@ import scipy.spatial
 @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