bug fix for training

Author: Davide-Miotti

pina/_src/solver/autoregressive_solver/autoregressive_solver.py

@@ -221,9 +221,12 @@ def loss_data(
         state_shape = unroll.shape[3:]
 
         losses = []
+        current_state = unroll[:, :, 0, ...]  # [B1, B2, *state_shape]
         for step in range(1, Twin):
 
-            model_input = self.preprocess_step(unroll, step=step, **kwargs)
+            model_input = self.preprocess_step(
+                current_state, unroll, step=step, **kwargs
+            )
             model_output = self.model_forward(model_input)
             predicted_state = self.postprocess_step(
                 model_output, unroll=unroll, step=step, **kwargs
@@ -233,6 +236,8 @@ def loss_data(
             step_loss = self._loss_fn(predicted_state, target_state, **kwargs)
             losses.append(step_loss)
 
+            current_state = predicted_state
+
         step_losses = torch.stack(losses)  # [unroll_length]
 
         with torch.no_grad():
@@ -252,21 +257,22 @@ def loss_data(
 
         return aggregation_strategy(step_losses * weights)
 
-    def preprocess_step(self, unroll, step=None, **kwargs):
+    def preprocess_step(self, current_state, unroll, step=None, **kwargs):
         """
         Pre-process the input unroll for the current step before feeding it to the model.
         This method can be overridden by subclasses to implement specific preprocessing logic.
 
-        :param torch.Tensor unroll: The unroll.
+        :param torch.Tensor current_state: The current state.
+        :param torch.Tensor unroll: The unroll, which can be used for context.
         :param int step: The current step index within the unroll.
         :kwargs: Additional keyword arguments for preprocessing.
         :return: The preprocessed unroll for the given step.
         :rtype: torch.Tensor
         """
         if step is None:
-            return unroll
+            return current_state
         else:
-            return unroll[:, :, step - 1, ...]
+            return current_state
 
     def model_forward(self, model_input, **kwargs):
         """
@@ -396,7 +402,7 @@ def predict(self, initial_input, num_steps, **kwargs):
         with torch.no_grad():
             for step in range(1, num_steps + 1):
                 model_input = self.preprocess_step(
-                    predictions[-1], step=None, **kwargs
+                    predictions[-1], unroll=predictions, step=None, **kwargs
                 )
                 next_state = self.model_forward(model_input)
                 next_state = self.postprocess_step(

pina/_src/solver/autoregressive_solver/autoregressive_solver_interface.py

@@ -97,17 +97,19 @@ def optimization_cycle(self, batch):
         return condition_loss
 
     @abstractmethod
-    def preprocess_step(self, tensor, step=None, **kwargs):
+    def preprocess_step(self, current_state, unroll, step=None, **kwargs):
         """
-        Preprocess the input state before passing it to the model.
+        Preprocess the current state before passing it to the model.
         This method can be overridden by subclasses to implement
         specific preprocessing steps.
 
-        :param torch.Tensor tensor: The tensor to be preprocessed.
+        :param torch.Tensor current_state: The current state to be preprocessed.
+        :param torch.Tensor unroll: The unroll tensor, which can be used for context.
+            During inference (predict), this may be a list of previous predictions.
         :param int step: The current step index within the unroll.
             By default is ``None``, which is meant to be used in inference.
         :kwargs: Additional keyword arguments for preprocessing.
-        :return: The preprocessed tensor for the given step.
+        :return: The preprocessed state for the given step.
         :rtype: torch.Tensor
         """
         pass