Unverified Commit 72a16b22 authored by Zihao Ye's avatar Zihao Ye Committed by GitHub
Browse files

[bugfix] Fix the behavior of select_topk (#2688)

* fix

* lint

* lint
parent e6bf54cd
...@@ -252,7 +252,10 @@ def select_topk(g, k, weight, nodes=None, edge_dir='in', ascending=False, ...@@ -252,7 +252,10 @@ def select_topk(g, k, weight, nodes=None, edge_dir='in', ascending=False,
""" """
# Rectify nodes to a dictionary # Rectify nodes to a dictionary
if nodes is None: if nodes is None:
nodes = {ntype: F.arange(0, g.number_of_nodes(ntype)) for ntype in g.ntypes} nodes = {
ntype: F.astype(F.arange(0, g.number_of_nodes(ntype)), g.idtype)
for ntype in g.ntypes
}
elif not isinstance(nodes, dict): elif not isinstance(nodes, dict):
if len(g.ntypes) > 1: if len(g.ntypes) > 1:
raise DGLError("Must specify node type when the graph is not homogeneous.") raise DGLError("Must specify node type when the graph is not homogeneous.")
......
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