Merge pull request #1002 from linzihan-backforward/master

FIX: bugs in ENMF
RUCAIBox · Oct 24, 2021 · 186a9f9 · 186a9f9
2 parents ec29419 + 0cc70e6
commit 186a9f9
Showing 1 changed file with 3 additions and 2 deletions.
diff --git a/recbole/model/general_recommender/enmf.py b/recbole/model/general_recommender/enmf.py
@@ -82,8 +82,9 @@ def calculate_loss(self, interaction):
                              self.item_embedding.weight.unsqueeze(1)).sum(dim=0)
 
         # shape: [embedding_size, embedding_size]
-        user_sum = torch.bmm(self.user_embedding.weight.unsqueeze(2),
-                             self.user_embedding.weight.unsqueeze(1)).sum(dim=0)
+        batch_user = self.user_embedding(user)
+        user_sum = torch.bmm(batch_user.unsqueeze(2),
+                             batch_user.unsqueeze(1)).sum(dim=0)
 
         # shape: [embedding_size, embedding_size]
         H_sum = torch.matmul(self.H_i.weight.t(), self.H_i.weight)