refactor: logging embeddings for both training and validation steps i…

…n multitas
IntelLabs · Jul 1, 2024 · d53cde4 · d53cde4
1 parent 43175b4
commit d53cde4
Showing 1 changed file with 6 additions and 0 deletions.
diff --git a/matsciml/models/base.py b/matsciml/models/base.py
@@ -3051,6 +3051,9 @@ def training_step(
             prog_bar=True,
             batch_size=batch_info["batch_size"],
         )
+        # optionally log embeddings
+        if self.hparams.log_embeddings and "embeddings" in batch:
+            self._log_embedding(batch["embeddings"])
         return losses
 
     def validation_step(
@@ -3105,6 +3108,9 @@ def validation_step(
             prog_bar=True,
             batch_size=batch_info["batch_size"],
         )
+        # optionally log embeddings
+        if self.hparams.log_embeddings and "embeddings" in batch:
+            self._log_embedding(batch["embeddings"])
         return losses
 
     @classmethod