Fix bug in test rmse computation, causing incorrect values (generally…

… inflated)
mpartio · Nov 7, 2023 · 6377d44 · 6377d44
1 parent 2d86715
commit 6377d44
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
@@ -338,7 +338,7 @@ def on_test_epoch_end(self):
  if self.trainer.is_global_zero:
  test_mae_rescaled = torch.mean(test_mae_tensor,
  dim=0) * self.data_std # (pred_steps, d_f)
- test_rmse_rescaled = torch.sqrt(torch.mean(test_mae_tensor,
+ test_rmse_rescaled = torch.sqrt(torch.mean(test_mse_tensor,
  dim=0)) * self.data_std # (pred_steps, d_f)
 
  mae_fig = vis.plot_error_map(test_mae_rescaled, step_length=self.step_length)