bugfix indexing batch-index for time

mllam · Jan 13, 2025 · 6672756 · 6672756
1 parent 7316a00
commit 6672756
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/neural_lam/models/ar_model.py b/neural_lam/models/ar_model.py
@@ -505,7 +505,7 @@ def plot_examples(self, batch, n_examples, split, prediction=None):
             prediction, target, _, _ = self.common_step(batch)
 
         target = batch[1]
-        time = batch[3]
+        time = batch[-1]
 
         # Rescale to original data scale
         prediction_rescaled = prediction * self.state_std + self.state_mean