apply time-indexing to support ar_steps_val > 1

leifdenby · Nov 29, 2024 · 52c4528 · 52c4528
1 parent 85160ce
commit 52c4528
Showing 1 changed file with 4 additions and 2 deletions.
diff --git a/neural_lam/models/ar_model.py b/neural_lam/models/ar_model.py
@@ -522,9 +522,11 @@ def plot_examples(self, batch, n_examples, split, prediction=None):
                         f"t={t_i} ({self._datastore.step_length * t_i} h)",
                         vrange=var_vrange,
                         da_prediction=da_prediction.isel(
-                            state_feature=var_i
+                            state_feature=var_i, time=t_i - 1
+                        ).squeeze(),
+                        da_target=da_target.isel(
+                            state_feature=var_i, time=t_i - 1
                         ).squeeze(),
-                        da_target=da_target.isel(state_feature=var_i).squeeze(),
                     )
                     for var_i, (var_name, var_unit, var_vrange) in enumerate(
                         zip(