FIX： Bugs in ENMF

RUCAIBox · Oct 14, 2021 · 0cc70e6 · 0cc70e6
1 parent b18820f
commit 0cc70e6
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)