log loss

discovery-unicamp · May 14, 2024 · add3be0 · add3be0
1 parent 77c0828
commit add3be0
Showing 1 changed file with 10 additions and 0 deletions.
diff --git a/minerva/models/nets/setr.py b/minerva/models/nets/setr.py
@@ -604,6 +604,16 @@ def _single_step(self, batch: torch.Tensor, batch_idx: int, step_name: str):
                 sync_dist=True,
             )
 
+        self.log(
+            f"{step_name}_loss",
+            loss,
+            on_step=True,
+            on_epoch=True,
+            prog_bar=True,
+            logger=True,
+            sync_dist=True,
+        )
+
         return loss
 
     def training_step(self, batch: torch.Tensor, batch_idx: int):