Skip to content

GitLab

Unverified Commit 8c28f982 authored by Muhammed Fatih BALIN's avatar Muhammed Fatih BALIN Committed by GitHub
Browse files

[GraphBolt] Slight `gb.expand_indptr` improvement on CPU. (#7189)

parent 290dadb9
Show whitespace changes
Inline Side-by-side
Showing with 1 addition and 1 deletion
graphbolt/src/expand_indptr.cc
View file @ 8c28f982
...@@ -23,7 +23,7 @@ torch::Tensor ExpandIndptr( ...@@ -23,7 +23,7 @@ torch::Tensor ExpandIndptr(
}); });
} }
if (!node_ids.has_value()) { if (!node_ids.has_value()) {
node_ids = torch::arange(indptr.size(0) - 1, indptr.options().dtype(dtype)); return torch::repeat_interleave(indptr.diff(), output_size).to(dtype);
} }
return node_ids.value().to(dtype).repeat_interleave( return node_ids.value().to(dtype).repeat_interleave(
indptr.diff(), 0, output_size); indptr.diff(), 0, output_size);
......
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