fix calculation of weighted loss

alibaba · Dec 16, 2024 · d7e5380 · d7e5380
1 parent 5097d1d
commit d7e5380
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 2 deletions.
diff --git a/tzrec/models/match_model.py b/tzrec/models/match_model.py
@@ -233,7 +233,9 @@ def _loss_impl(
             label = _zero_int_label(pred)
         losses[loss_name] = self._loss_modules[loss_name](pred, label)
         if self._sample_weight:
-            losses[loss_name] = torch.mean(losses[loss_name] * sample_weight)
+            losses[loss_name] = torch.sum(
+                losses[loss_name] * sample_weight
+            ) / torch.sum(sample_weight)
 
         return losses
 

diff --git a/tzrec/models/rank_model.py b/tzrec/models/rank_model.py
@@ -215,7 +215,9 @@ def _loss_impl(
             raise ValueError(f"loss[{loss_type}] is not supported yet.")
 
         if sample_weight_name:
-            losses[loss_name] = torch.mean(losses[loss_name] * sample_weights)
+            losses[loss_name] = torch.sum(
+                losses[loss_name] * sample_weights
+            ) / torch.sum(sample_weights)
         return losses
 
     def loss(