"vscode:/vscode.git/clone" did not exist on "7b07e123067bbc49cadd77270e8d898c86f71aaa"
Unverified Commit 7b4c8c77 authored by Mingbang Wang's avatar Mingbang Wang Committed by GitHub
Browse files

[GraphBolt] Update `_init_all_nodes_set` (#7168)

parent 1e589e85
......@@ -828,11 +828,18 @@ class OnDiskDataset(Dataset):
)
return None
num_nodes = graph.num_nodes
dtype = graph.indices.dtype
if isinstance(num_nodes, int):
return ItemSet(num_nodes, names="seed_nodes")
return ItemSet(
torch.tensor(num_nodes, dtype=dtype),
names="seed_nodes",
)
else:
data = {
node_type: ItemSet(num_node, names="seed_nodes")
node_type: ItemSet(
torch.tensor(num_node, dtype=dtype),
names="seed_nodes",
)
for node_type, num_node in num_nodes.items()
}
return ItemSetDict(data)
......
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