[Fix] fetch negative entities
@@ -106,7 +106,9 @@ class KGEModel(nn.Layer):
pos_h = paddle.unsqueeze(ent_emb[h], axis=1)
pos_t = paddle.unsqueeze(ent_emb[t], axis=1)
pos_r = paddle.unsqueeze(pos_r, axis=1)
- neg_ents = ent_emb[neg_ents]
+ neg_ents_shape = neg_ents.shape
+ neg_ents = ent_emb[paddle.reshape(neg_ents, (-1, ))]
+ neg_ents = paddle.reshape(neg_ents, [*neg_ents_shape, -1])
pos_score = self._score_func(pos_h, pos_r, pos_t)
if neg_mode == 'tail':
