refactor: making validation and test log to progress bar as well

laserkelvin · Jul 19, 2024 · f19fa79 · f19fa79
1 parent 17c69c2
commit f19fa79
Showing 1 changed file with 4 additions and 4 deletions.
diff --git a/matsciml/models/base.py b/matsciml/models/base.py
@@ -1083,7 +1083,7 @@ def validation_step(
                 "Unable to parse batch size from data, defaulting to `None` for logging.",
             )
             batch_size = None
-        self.log_dict(metrics, batch_size=batch_size)
+        self.log_dict(metrics, batch_size=batch_size, on_step=True, prog_bar=True)
         if self.hparams.log_embeddings and "embeddings" in batch:
             self._log_embedding(batch["embeddings"])
         return loss_dict
@@ -1105,7 +1105,7 @@ def test_step(
                 "Unable to parse batch size from data, defaulting to `None` for logging.",
             )
             batch_size = None
-        self.log_dict(metrics, batch_size=batch_size)
+        self.log_dict(metrics, batch_size=batch_size, on_epoch=True, prog_bar=True)
         if self.hparams.log_embeddings and "embeddings" in batch:
             self._log_embedding(batch["embeddings"])
         return loss_dict
@@ -1510,7 +1510,7 @@ def validation_step(
                 "Unable to parse batch size from data, defaulting to `None` for logging."
             )
             batch_size = None
-        self.log_dict(metrics, batch_size=batch_size)
+        self.log_dict(metrics, batch_size=batch_size, on_step=True, prog_bar=True)
         return loss_dict
 
     def test_step(
@@ -1531,7 +1531,7 @@ def test_step(
                 "Unable to parse batch size from data, defaulting to `None` for logging."
             )
             batch_size = None
-        self.log_dict(metrics, batch_size=batch_size)
+        self.log_dict(metrics, batch_size=batch_size, on_step=True, prog_bar=True)
         return loss_dict
 
     def on_train_batch_start(self, batch: Any, batch_idx: int) -> Optional[int]: