From 2d39262198b8286c65a944cb96699ddc6ac65bcb Mon Sep 17 00:00:00 2001 From: Cathal OBrien Date: Thu, 21 Nov 2024 15:58:15 +0000 Subject: [PATCH] remove prog fields variable we can save 2GB of memory at 9km by removing this variable --- src/anemoi/inference/runner.py | 5 +---- 1 file changed, 1 insertion(+), 4 deletions(-) diff --git a/src/anemoi/inference/runner.py b/src/anemoi/inference/runner.py index e9cbd27..93e8bf3 100644 --- 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