Merge pull request #128 from Sujit-O/development

fix bug in trainer.py projection train step.
Sujit-O · Apr 10, 2020 · 561d4ea · 561d4ea
2 parents 9e5b2de + 110cc01
commit 561d4ea
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/pykg2vec/utils/trainer.py b/pykg2vec/utils/trainer.py
@@ -174,7 +174,7 @@ def train_step_projection(self, h, r, t, hr_t, tr_h):
                     tr_h = tr_h * (1.0 - self.config.label_smoothing) + 1.0 / self.config.kg_meta.tot_entity
 
                 pred_tails = self.model.forward(h, r, direction="tail") # (h, r) -> hr_t forward
-                pred_heads = self.model.forward(h, r, direction="head") # (t, r) -> tr_h backward
+                pred_heads = self.model.forward(t, r, direction="head") # (t, r) -> tr_h backward
 
                 loss_tails = tf.reduce_mean(tf.keras.backend.binary_crossentropy(hr_t, pred_tails))
                 loss_heads = tf.reduce_mean(tf.keras.backend.binary_crossentropy(tr_h, pred_heads))
@@ -183,7 +183,7 @@ def train_step_projection(self, h, r, t, hr_t, tr_h):
 
             else:
                 loss_tails = self.model.forward(h, r, hr_t, direction="tail") # (h, r) -> hr_t forward
-                loss_heads = self.model.forward(h, r, tr_h, direction="head") # (t, r) -> tr_h backward
+                loss_heads = self.model.forward(t, r, tr_h, direction="head") # (t, r) -> tr_h backward
 
                 loss = loss_tails + loss_heads