Fix wrong loss

lightly-ai · Sep 16, 2023 · e8c5d7d · e8c5d7d
1 parent eb1b784
commit e8c5d7d
Showing 1 changed file with 5 additions and 5 deletions.
diff --git a/lightly/loss/tico_loss.py b/lightly/loss/tico_loss.py
@@ -109,11 +109,11 @@ def forward(
 
         # compute loss
         C = self.beta * self.C + (1 - self.beta) * B
-        loss = (
-            1
-            - (z_a * z_b).sum(dim=1).mean()
-            + self.rho * (torch.mm(z_a, C) * z_a).sum(dim=1).mean()
-        )
+
+        transformative_invariance_loss = -(z_a * z_b).sum(dim=1).mean()
+        covariance_contrast_loss = self.rho * (torch.mm(z_a, C) * z_a).sum(dim=1).mean()
+
+        loss = transformative_invariance_loss + covariance_contrast_loss
 
         # update covariance matrix
         if update_covariance_matrix: