"...git@developer.sourcefind.cn:renzhc/diffusers_dcu.git" did not exist on "b33bd91faea027ebe744a90c36fa92610899f822"
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): ...@@ -828,11 +828,18 @@ class OnDiskDataset(Dataset):
) )
return None return None
num_nodes = graph.num_nodes num_nodes = graph.num_nodes
dtype = graph.indices.dtype
if isinstance(num_nodes, int): if isinstance(num_nodes, int):
return ItemSet(num_nodes, names="seed_nodes") return ItemSet(
torch.tensor(num_nodes, dtype=dtype),
names="seed_nodes",
)
else: else:
data = { 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() for node_type, num_node in num_nodes.items()
} }
return ItemSetDict(data) 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