Merge pull request #1873 from TayTroye/oks

FIX(trainer.py):calculate_loss error when use multi-gpu
RUCAIBox · Sep 21, 2023 · 96eb311 · 96eb311
2 parents f9104dc + b66f079
commit 96eb311
Showing 1 changed file with 5 additions and 1 deletion.
diff --git a/recbole/model/general_recommender/pop.py b/recbole/model/general_recommender/pop.py
@@ -6,6 +6,10 @@
 # @Time   : 2020/11/9
 # @Author : Zihan Lin
 # @Email  : [email protected]
+# UPDATE
+# @Time   :2023/9/21
+# @Author : Kesha Ou
+# @Email  :[email protected]
 
 r"""
 Pop
@@ -43,7 +47,7 @@ def calculate_loss(self, interaction):
 
         self.max_cnt = torch.max(self.item_cnt, dim=0)[0]
 
-        return torch.nn.Parameter(torch.zeros(1))
+        return torch.nn.Parameter(torch.zeros(1)).to(self.device)
 
     def predict(self, interaction):
         item = interaction[self.ITEM_ID]