Skip to content

GitLab

Unverified Commit 6566c31f authored by Chang Liu's avatar Chang Liu Committed by GitHub
Browse files

Fix ogb/ogbn-mag/heter-RGCN example (#4839) 


Co-authored-by: default avatarMufei Li <mufeili1996@gmail.com>
parent 344be1ef
Show whitespace changes
Inline Side-by-side
Showing with 5 additions and 2 deletions
examples/pytorch/ogb/ogbn-mag/hetero_rgcn.py
View file @ 6566c31f
...@@ -266,11 +266,13 @@ def train( ...@@ -266,11 +266,13 @@ def train(
category category
] # we only predict the nodes with type "category" ] # we only predict the nodes with type "category"
batch_size = seeds.shape[0] batch_size = seeds.shape[0]
input_nodes_indexes = input_nodes["paper"].to(g.device)
seeds = seeds.to(labels.device)
emb = extract_embed(node_embed, input_nodes) emb = extract_embed(node_embed, input_nodes)
# Add the batch's raw "paper" features # Add the batch's raw "paper" features
emb.update( emb.update(
{"paper": g.ndata["feat"]["paper"][input_nodes["paper"]]} {"paper": g.ndata["feat"]["paper"][input_nodes_indexes]}
) )
emb = {k: e.to(device) for k, e in emb.items()} emb = {k: e.to(device) for k, e in emb.items()}
...@@ -334,10 +336,11 @@ def test(g, model, node_embed, y_true, device, split_idx): ...@@ -334,10 +336,11 @@ def test(g, model, node_embed, y_true, device, split_idx):
category category
] # we only predict the nodes with type "category" ] # we only predict the nodes with type "category"
batch_size = seeds.shape[0] batch_size = seeds.shape[0]
input_nodes_indexes = input_nodes["paper"].to(g.device)
emb = extract_embed(node_embed, input_nodes) emb = extract_embed(node_embed, input_nodes)
# Get the batch's raw "paper" features # Get the batch's raw "paper" features
emb.update({"paper": g.ndata["feat"]["paper"][input_nodes["paper"]]}) emb.update({"paper": g.ndata["feat"]["paper"][input_nodes_indexes]})
emb = {k: e.to(device) for k, e in emb.items()} emb = {k: e.to(device) for k, e in emb.items()}
logits = model(emb, blocks)[category] logits = model(emb, blocks)[category]
......
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