remove prog fields variable

we can save 2GB of memory at 9km by removing this variable
ecmwf · Nov 21, 2024 · 2d39262 · 2d39262
1 parent 9c7b77c
commit 2d39262
Showing 1 changed file with 1 addition and 4 deletions.
diff --git a/src/anemoi/inference/runner.py b/src/anemoi/inference/runner.py
@@ -493,9 +493,6 @@ def get_most_recent_datetime(input_fields):
                         )
 
             # Next step
-
-            prognostic_fields = y_pred[..., prognostic_output_mask]
-
             # Compute new forcing
 
             forcing = forcing_and_constants(
@@ -508,7 +505,7 @@ def get_most_recent_datetime(input_fields):
 
             # Update dynamic tensor for next iteration
             input_tensor_torch = input_tensor_torch.roll(-1, dims=1)
-            input_tensor_torch[:, -1, :, prognostic_input_mask] = prognostic_fields
+            input_tensor_torch[:, -1, :, prognostic_input_mask] = y_pred[..., prognostic_output_mask]
             if computed_forcing_mask:
                 input_tensor_torch[:, -1, :, computed_forcing_mask] = forcing