From 52c452879f56c7f982cfd5d55a5259f37cb6b030 Mon Sep 17 00:00:00 2001 From: Leif Denby Date: Fri, 29 Nov 2024 15:05:36 +0100 Subject: [PATCH] apply time-indexing to support ar_steps_val > 1 --- neural_lam/models/ar_model.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/neural_lam/models/ar_model.py b/neural_lam/models/ar_model.py index 0d8e6e3c..44baf9c2 100644 --- 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(