Allow for config of expand_keys

ecmwf · Oct 28, 2024 · 6e38c9e · 6e38c9e
1 parent 745b812
commit 6e38c9e
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 1 deletion.
diff --git a/src/anemoi/training/config/diagnostics/eval_rollout.yaml b/src/anemoi/training/config/diagnostics/eval_rollout.yaml
@@ -94,6 +94,8 @@ log:
     terminal: True
     run_name: null # If set to null, the run name will be the a random UUID
     on_resume_create_child: True
+    expand_hyperparams: # Which keys in hyperparams to expand
+      - config
   interval: 100
 
 enable_progress_bar: True

diff --git a/src/anemoi/training/diagnostics/logger.py b/src/anemoi/training/diagnostics/logger.py
@@ -73,7 +73,10 @@ def get_mlflow_logger(config: DictConfig) -> None:
     )
     config_params = OmegaConf.to_container(config, resolve=True)
 
-    logger.log_hyperparams(config_params, expand_keys=["config"])
+    logger.log_hyperparams(
+        config_params,
+        expand_keys=config.diagnostics.log.mlflow.get("expand_hyperparams", ["config"]),
+    )
 
     if config.diagnostics.log.mlflow.terminal:
         logger.log_terminal_output(artifact_save_dir=config.hardware.paths.plots)