Update neural_lam/models/ar_model.py

Co-authored-by: Joel Oskarsson <[email protected]>
mllam · Sep 27, 2024 · 8453c2b · 8453c2b
1 parent aee0b1c
commit 8453c2b
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
@@ -291,7 +291,7 @@ def validation_step(self, batch, batch_idx):
         val_log_dict = {
             f"val_loss_unroll{step}": time_step_loss[step - 1]
             for step in self.args.val_steps_to_log
-            if step < len(time_step_loss)
+            if step <= len(time_step_loss)
         }
         val_log_dict["val_mean_loss"] = mean_loss
         self.log_dict(