Wrapping pytorch autolog with rank_zero_only decorator to avoid multi…

…ple runs during multi gpu training

Signed-off-by: Shrinath Suresh <[email protected]>

mlflow/pytorch/__init__.py

@@ -54,6 +54,7 @@
 )
 from mlflow.tracking._model_registry import DEFAULT_AWAIT_MAX_SLEEP_SECONDS
 from mlflow.utils.autologging_utils import autologging_integration, safe_patch
+from pytorch_lightning.utilities import rank_zero_only
 
 FLAVOR_NAME = "pytorch"
 
@@ -878,6 +879,7 @@ def load_state_dict(state_dict_uri, **kwargs):
 
 
 @autologging_integration(FLAVOR_NAME)
+@rank_zero_only
 def autolog(
     log_every_n_epoch=1,
     log_every_n_step=None,