fix neg sample eval bugs of diffrec and ldiffrec

RUCAIBox · Feb 27, 2024 · 77b7cb6 · 77b7cb6
1 parent 77a60ae
commit 77b7cb6
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