Merge pull request #27 from lightonai/constrastive-loss-label

add-labels
lightonai · Aug 8, 2024 · f2c3e83 · f2c3e83
2 parents da24370 + 0d3ec82
commit f2c3e83
Showing 1 changed file with 6 additions and 2 deletions.
diff --git a/giga_cherche/losses/contrastive.py b/giga_cherche/losses/contrastive.py
@@ -120,7 +120,9 @@ def __init__(
         self.size_average = size_average
 
     def forward(
-        self, sentence_features: Iterable[dict[str, Tensor]], **kwargs
+        self,
+        sentence_features: Iterable[dict[str, Tensor]],
+        labels: torch.Tensor | None = None,
     ) -> torch.Tensor:
         """Compute the Constrastive loss.
 
@@ -157,5 +159,7 @@ def forward(
         # compute constrastive loss using cross-entropy over the distances
 
         return F.cross_entropy(
-            distances, labels, reduction="mean" if self.size_average else "sum"
+            input=distances,
+            target=labels,
+            reduction="mean" if self.size_average else "sum",
         )