Rename AIFSMLflowLogger -> AnemoiMLflowLogger

ecmwf · Jul 17, 2024 · edcd1e1 · edcd1e1
1 parent 290e48f
commit edcd1e1
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 3 deletions.
diff --git a/src/anemoi/training/diagnostics/callbacks.py b/src/anemoi/training/diagnostics/callbacks.py
@@ -540,9 +540,9 @@ def tracker_metadata(self, trainer):
         elif self.config.diagnostics.log.mlflow.enabled:
             self._tracker_name = "mlflow"
 
-            from anemoi.training.diagnostics.mlflow.logger import AIFSMLflowLogger
+            from anemoi.training.diagnostics.mlflow.logger import AnemoiMLflowLogger
 
-            mlflow_logger = next(logger for logger in trainer.loggers if isinstance(logger, AIFSMLflowLogger))
+            mlflow_logger = next(logger for logger in trainer.loggers if isinstance(logger, AnemoiMLflowLogger))
             run_id = mlflow_logger.run_id
             run = mlflow_logger._mlflow_client.get_run(run_id)
 

diff --git a/src/anemoi/training/diagnostics/mlflow/logger.py b/src/anemoi/training/diagnostics/mlflow/logger.py
@@ -265,7 +265,7 @@ def _remove_csi(line):
         self.experiment.log_artifact(self.run_id, str(self.file_save_path))
 
 
-class AIFSMLflowLogger(MLFlowLogger):
+class AnemoiMLflowLogger(MLFlowLogger):
     """A custom MLflow logger that logs terminal output."""
 
     def __init__(