Update metadata handling

ecmwf · May 28, 2024 · c3ab66c · c3ab66c
1 parent 6a00010
commit c3ab66c
Showing 1 changed file with 2 additions and 7 deletions.
diff --git a/src/anemoi/inference/runner.py b/src/anemoi/inference/runner.py
@@ -66,6 +66,7 @@ def run(
         output_callback=ignore,
         autocast=None,
         progress_callback=ignore,
+        add_ensemble_dimension=False,
     ):
 
         if autocast is None:
@@ -287,13 +288,7 @@ def wrapper(x):
 
             return wrapper
 
-        if True:  # self.add_ensemble_dimension:
-            LOGGER.warning("🚨" * 80)
-            LOGGER.warning("Adding ensemble dimension.")
-            LOGGER.warning("If you are using that flags, your are using unsupported code")
-            LOGGER.warning("that can be removed any time in the near future")
-            LOGGER.warning("🚨" * 80)
-
+        if add_ensemble_dimension:
             model.predict_step = add_ensemble_dim(model.predict_step)
 
         prognostic_params = self.checkpoint.prognostic_params