From 0eab1e42b24db5b6006c03954435118f0fcbd3e9 Mon Sep 17 00:00:00 2001 From: festeh Date: Thu, 10 Oct 2019 00:47:17 +0300 Subject: [PATCH] add tags argument to MLFlowLogger (#349) --- pytorch_lightning/logging/mlflow_logger.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/pytorch_lightning/logging/mlflow_logger.py b/pytorch_lightning/logging/mlflow_logger.py index 970dfe0309..8fb79f5366 100644 --- a/pytorch_lightning/logging/mlflow_logger.py +++ b/pytorch_lightning/logging/mlflow_logger.py @@ -9,11 +9,12 @@ logger = getLogger(__name__) class MLFlowLogger(LightningLoggerBase): - def __init__(self, experiment_name, tracking_uri=None): + def __init__(self, experiment_name, tracking_uri=None, tags=None): super().__init__() self.client = mlflow.tracking.MlflowClient(tracking_uri) self.experiment_name = experiment_name self._run_id = None + self.tags = tags @property def run_id(self): @@ -28,7 +29,7 @@ class MLFlowLogger(LightningLoggerBase): self.client.create_experiment(self.experiment_name) experiment = self.client.get_experiment_by_name(self.experiment_name) - run = self.client.create_run(experiment.experiment_id) + run = self.client.create_run(experiment.experiment_id, tags=self.tags) self._run_id = run.info.run_id return self._run_id