remove artifacts from earlier merging/rebase

mllam · Dec 11, 2024 · ae69f3f · ae69f3f
1 parent b5ebe6f
commit ae69f3f
Showing 1 changed file with 1 addition and 2 deletions.
diff --git a/neural_lam/models/ar_model.py b/neural_lam/models/ar_model.py
@@ -95,7 +95,6 @@ def __init__(
             # Store constant per-variable std.-dev. weighting
             # NOTE that this is the inverse of the multiplicative weighting
             # in wMSE/wMAE
-            # TODO: Do we need param_weights for this?
             self.register_buffer(
                 "per_var_std",
                 self.diff_std / torch.sqrt(self.feature_weights),
@@ -262,7 +261,7 @@ def unroll_prediction(self, init_states, forcing_features, true_states):
                 pred_std_list, dim=1
             )  # (B, pred_steps, num_grid_nodes, d_f)
         else:
-            pred_std = self.diff_std  # (d_f,)
+            pred_std = self.per_var_std  # (d_f,)
 
         return prediction, pred_std