fix link prediction (dmlc#3485)

niushixiong · Nov 9, 2021 · f7360c3 · f7360c3
1 parent db78fac
commit f7360c3
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 2 deletions.
diff --git a/tutorials/blitz/4_link_predict.py b/tutorials/blitz/4_link_predict.py
@@ -99,7 +99,7 @@
 
 neg_eids = np.random.choice(len(neg_u), g.number_of_edges())
 test_neg_u, test_neg_v = neg_u[neg_eids[:test_size]], neg_v[neg_eids[:test_size]]
-train_neg_u, train_neg_v = neg_u[neg_eids[train_size:]], neg_v[neg_eids[train_size:]]
+train_neg_u, train_neg_v = neg_u[neg_eids[test_size:]], neg_v[neg_eids[test_size:]]
 
 
 ######################################################################

diff --git a/tutorials/large/L2_large_link_prediction.py b/tutorials/large/L2_large_link_prediction.py
@@ -395,7 +395,7 @@ def closure():
 # You also need to label the edges, 1 if positive and 0 if negative.
 #
 
-test_src = torch.cat([test_pos_src, test_neg_src])
+test_src = torch.cat([test_pos_src, test_pos_dst])
 test_dst = torch.cat([test_neg_src, test_neg_dst])
 test_graph = dgl.graph((test_src, test_dst), num_nodes=graph.num_nodes())
 test_graph.edata['label'] = torch.cat(