fix typo (missing datastore) in ARModel.on_test_epoch_end

mllam · Nov 13, 2024 · df4d39c · df4d39c
1 parent f97719b
commit df4d39c
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/neural_lam/models/ar_model.py b/neural_lam/models/ar_model.py
@@ -601,7 +601,8 @@ def on_test_epoch_end(self):
                 vis.plot_spatial_error(
                     error=loss_map,
                     datastore=self._datastore,
-                    title=f"Test loss, t={t_i} ({self.step_length * t_i} h)",
+                    title=f"Test loss, t={t_i} "
+                    f"({self._datastore.step_length * t_i} h)",
                 )
                 for t_i, loss_map in zip(
                     self.args.val_steps_to_log, mean_spatial_loss