[NeMo-UX] log val loss (NVIDIA#9814) (NVIDIA#9831)

Signed-off-by: ashors1 <[email protected]> Co-authored-by: Anna Shors <[email protected]> Signed-off-by: Vivian Chen <[email protected]>
xuanzic · Aug 1, 2024 · e12d357 · e12d357
1 parent d2a0b46
commit e12d357
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/nemo/lightning/pytorch/strategies.py b/nemo/lightning/pytorch/strategies.py
@@ -441,7 +441,9 @@ def validation_step(self, dataloader_iter, *args: Any, **kwargs: Any) -> STEP_OU
         kwargs = self._update_step_kwargs(dataloader_iter, kwargs, "validation")
 
         with self.precision_plugin.val_step_context():  # TODO: Do we need this?
-            return self.model(dataloader_iter, forward_only=True, *args, **kwargs)
+            out = self.model(dataloader_iter, forward_only=True, *args, **kwargs)
+            self.lightning_module.log('val_loss', out, rank_zero_only=True, batch_size=1)
+            return out
 
     @override
     def test_step(self, dataloader_iter, *args: Any, **kwargs: Any) -> STEP_OUTPUT: