Adding rank zero decorator to patched_fit method - to log model only …

…once in multi gpu environment Signed-off-by: Shrinath Suresh <[email protected]>
drsantos89 · May 9, 2022 · 97ca87e · 97ca87e
1 parent c577c7f
commit 97ca87e
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/mlflow/pytorch/_pytorch_autolog.py b/mlflow/pytorch/_pytorch_autolog.py
@@ -325,6 +325,7 @@ def _log_early_stop_metrics(early_stop_callback, client, run_id):
     client.log_metrics(run_id, metrics)
 
 
+@rank_zero_only
 def patched_fit(original, self, *args, **kwargs):
     """
     A patched implementation of `pytorch_lightning.Trainer.fit` which enables logging the