Stitcher: Console log only on rank0

neuro-galaxy · Nov 16, 2024 · 6308be5 · 6308be5
1 parent 024cd27
commit 6308be5
Showing 1 changed file with 4 additions and 4 deletions.
diff --git a/torch_brain/utils/stitcher.py b/torch_brain/utils/stitcher.py
@@ -185,18 +185,18 @@ def on_validation_epoch_end(self, trainer, pl_module, prefix="val"):
 
         # log the metrics
         self.log_dict(metrics)
-        if not self.quiet:
-            logging.info(f"Logged {len(metrics)} {prefix} metrics.")
 
         metrics_data = []
         for metric_name, metric_value in metrics.items():
             metrics_data.append({"metric": metric_name, "value": metric_value.item()})
 
         metrics_df = pd.DataFrame(metrics_data)
-        if not self.quiet:
-            rprint(metrics_df)
 
         if trainer.is_global_zero:
+            if not self.quiet:
+                logging.info(f"Logged {len(metrics)} {prefix} metrics.")
+                rprint(metrics_df)
+
             for logger in trainer.loggers:
                 if isinstance(logger, L.pytorch.loggers.TensorBoardLogger):
                     logger.experiment.add_text(