From f94e46c59f3c1fa9481f0aa754069be842c9cc83 Mon Sep 17 00:00:00 2001
From: Jakob Schloer <jakob.schloer@gmail.com>
Date: Tue, 27 Aug 2024 14:31:27 +0000
Subject: [PATCH] Move residual from EncDecProcessor to Interface.

---
 src/anemoi/models/interface/__init__.py              | 12 ++++++++++--
 .../models/models/encoder_processor_decoder.py       |  2 --
 2 files changed, 10 insertions(+), 4 deletions(-)

diff --git a/src/anemoi/models/interface/__init__.py b/src/anemoi/models/interface/__init__.py
index d984c8a..045b966 100644
--- a/src/anemoi/models/interface/__init__.py
+++ b/src/anemoi/models/interface/__init__.py
@@ -13,6 +13,7 @@
 import torch
 from anemoi.utils.config import DotDict
 from hydra.utils import instantiate
+from torch.distributed.distributed_c10d import ProcessGroup
 from torch_geometric.data import HeteroData
 
 from anemoi.models.models.encoder_processor_decoder import AnemoiModelEncProcDec
@@ -96,8 +97,15 @@ def _build_model(self) -> None:
             config=self.config, data_indices=self.data_indices, graph_data=self.graph_data
         )
 
-        # Use the forward method of the model directly
-        self.forward = self.model.forward
+    def forward(self, x: torch.Tensor, model_comm_group: Optional[ProcessGroup] = None) -> torch.Tensor:
+        if self.tendency_mode:
+            # Predict tendency
+            x_pred = self.model.forward(x, model_comm_group)
+        else:
+            # Predict state by adding residual connection (just for the prognostic variables)
+            x_pred = self.model.forward(x, model_comm_group)
+            x_pred[..., self.model._internal_output_idx] += x[:, -1, :, :, self.model._internal_input_idx]
+        return x_pred
 
     def predict_step(self, batch: torch.Tensor) -> torch.Tensor:
         """Prediction step for the model.
diff --git a/src/anemoi/models/models/encoder_processor_decoder.py b/src/anemoi/models/models/encoder_processor_decoder.py
index 0f37474..5cda6e1 100644
--- a/src/anemoi/models/models/encoder_processor_decoder.py
+++ b/src/anemoi/models/models/encoder_processor_decoder.py
@@ -248,6 +248,4 @@ def forward(self, x: Tensor, model_comm_group: Optional[ProcessGroup] = None) ->
             .clone()
         )
 
-        # residual connection (just for the prognostic variables)
-        x_out[..., self._internal_output_idx] += x[:, -1, :, :, self._internal_input_idx]
         return x_out