Fix loss logging

berenslab · Dec 13, 2024 · 817fb28 · 817fb28
1 parent e242f47
commit 817fb28
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 4 deletions.
diff --git a/tsimcne/losses/infonce.py b/tsimcne/losses/infonce.py
@@ -83,7 +83,11 @@ def forward(self, features, backbone_features=None, labels=None):
         raw_uniformity = logsumexp_1 + logsumexp_2
 
         loss = -(self.exaggeration * tempered_alignment - raw_uniformity / 2)
-        return loss
+        return dict(
+            loss=loss,
+            ta=-tempered_alignment,
+            ru=raw_uniformity / 2,
+        )
 
 
 class InfoNCEGaussian(InfoNCECauchy):
@@ -111,7 +115,11 @@ def forward(self, features, backbone_features=None, labels=None):
         raw_uniformity = logsumexp_1 + logsumexp_2
 
         loss = -(tempered_alignment - raw_uniformity / 2)
-        return loss
+        return dict(
+            loss=loss,
+            ta=-tempered_alignment,
+            ru=raw_uniformity / 2,
+        )
 
 
 class InfoNCELoss(LossBase):

diff --git a/tsimcne/tsimcne.py b/tsimcne/tsimcne.py
@@ -356,9 +356,9 @@ def validation_step(self, batch, batch_idx, dataloader_idx=0):
             features, backbone_features = self(batch)
 
             # backbone_features are unused in infonce loss
-            loss = self.loss(features)
+            lossd = self.loss(features)
 
-            return dict(loss=loss)
+            return lossd
 
         elif dataloader_idx == 1:
             features, backbone_features = self(batch)