Skip to content

GitLab

"...git@developer.sourcefind.cn:OpenDAS/megatron-lm.git" did not exist on "688448db7547be90203440cfd105703d8a853f39"
Commit aff91e0e authored by rusty1s's avatar rusty1s
Browse files

fix jit bug

parent 6a8502cb
Hide whitespace changes
Inline Side-by-side
Showing with 1 addition and 1 deletion
torch_cluster/radius.py
View file @ aff91e0e
...@@ -7,7 +7,7 @@ import scipy.spatial ...@@ -7,7 +7,7 @@ import scipy.spatial
@torch.jit.script @torch.jit.script
def sample(col: torch.Tensor, count: int) -> torch.Tensor: def sample(col: torch.Tensor, count: int) -> torch.Tensor:
if col.size(0) > count: 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 return col
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment