Merge pull request #1999 from BishopLiu/master

FIX: fix neg sample eval bugs for diffrec and ldiffrec
RUCAIBox · Feb 27, 2024 · f95fde5 · f95fde5
2 parents de6690e + f11c765
commit f95fde5
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 2 deletions.
diff --git a/recbole/model/general_recommender/diffrec.py b/recbole/model/general_recommender/diffrec.py
@@ -328,7 +328,7 @@ def full_sort_predict(self, interaction):
     def predict(self, interaction):
         item = interaction[self.ITEM_ID]
         x_t = self.full_sort_predict(interaction)
-        scores = x_t[:, item]
+        scores = x_t[torch.arange(len(item)).to(self.device), item]
         return scores
 
     def calculate_loss(self, interaction):

diff --git a/recbole/model/general_recommender/ldiffrec.py b/recbole/model/general_recommender/ldiffrec.py
@@ -335,7 +335,7 @@ def full_sort_predict(self, interaction):
     def predict(self, interaction):
         item = interaction[self.ITEM_ID]
         x_t = self.full_sort_predict(interaction)
-        scores = x_t[:, item]
+        scores = x_t[torch.arange(len(item)).to(self.device), item]
         return scores