From 4a790172bc48512e341c5942f36ebf422d0d35d1 Mon Sep 17 00:00:00 2001 From: Zichen Wang Date: Mon, 27 Nov 2023 23:08:50 +0000 Subject: [PATCH] clean-ups --- python/graphstorm/model/edge_decoder.py | 6 +----- 1 file changed, 1 insertion(+), 5 deletions(-) diff --git a/python/graphstorm/model/edge_decoder.py b/python/graphstorm/model/edge_decoder.py index 335911b11b..43b4c5df98 100644 --- a/python/graphstorm/model/edge_decoder.py +++ b/python/graphstorm/model/edge_decoder.py @@ -724,11 +724,7 @@ def calc_retrieval_scores(self, emb, pos_pairs, device): scores = {} pos_scores = calc_dot_pos_score(pos_src_emb, pos_dst_emb) neg_dst_emb = emb[vtype][np.arange(emb[vtype].shape[0])].to(device) - # neg_dst_emb should contains train nodes only: - # v_train_mask = g.nodes[vtype].data['train_mask'][np.arange(g.number_of_nodes(vtype))] - # train_nids = np.where(v_train_mask)[0] - # neg_dst_emb = emb[vtype][train_nids].to(device) - neg_scores = th.mm(pos_src_emb, neg_dst_emb.transpose(0, 1)) # [n_pos, n_train] + neg_scores = th.mm(pos_src_emb, neg_dst_emb.transpose(0, 1)) # [n_pos, n_embs] # gloo with cpu will consume less GPU memory neg_scores = neg_scores.cpu() \ if is_distributed() and get_backend() == "gloo" \