use progress bar

ecmwf · May 25, 2024 · 8e9d712 · 8e9d712
1 parent 8f2f6ea
commit 8e9d712
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 5 deletions.
diff --git a/pyproject.toml b/pyproject.toml
@@ -39,7 +39,7 @@ classifiers = [
     "Operating System :: OS Independent",
 ]
 
-dependencies = ["anemoi-utils"]
+dependencies = ["anemoi-utils", "semantic-version"]
 
 [project.optional-dependencies]
 

diff --git a/src/anemoi/inference/runner.py b/src/anemoi/inference/runner.py
@@ -12,6 +12,7 @@
 
 import numpy as np
 import torch
+from anemoi.utils.timer import Timer
 
 from .checkpoint import Checkpoint
 
@@ -224,8 +225,7 @@ def run(
                     )
                     raise ValueError(f"Field '{name}' has NaNs and is not marked as imputable")
 
-        # with self.timer(f"Loading {self.checkpoint}"):
-        if True:
+        with Timer(f"Loading {self.checkpoint}"):
             try:
                 model = torch.load(
                     self.checkpoint.path,
@@ -300,7 +300,7 @@ def wrapper(x):
 
         # with self.stepper(self.hour_steps) as stepper:
 
-        for i in range(lead_time // self.hour_steps):
+        for i in progress_callback(range(lead_time // self.hour_steps)):
             step = (i + 1) * self.hour_steps
 
             # Predict next state of atmosphere
@@ -364,7 +364,7 @@ def wrapper(x):
             input_tensor_torch[:, -1, :, prognostic_input_mask] = prognostic_fields
             input_tensor_torch[:, -1, :, computed_forcing_mask] = forcing
 
-            progress_callback(i)
+            # progress_callback(i)
 
     @cached_property
     def hour_steps(self):