Adjust forcing dimensionalities after fix

mllam · Dec 5, 2024 · 29063c8 · 29063c8
1 parent a0f229b
commit 29063c8
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/neural_lam/models/ar_model.py b/neural_lam/models/ar_model.py
@@ -109,7 +109,7 @@ def __init__(
             2 * self.grid_output_dim
             + grid_static_dim
             # Temporal Embedding counts as one additional forcing_feature
-            + 2*(num_forcing_vars)
+            + (num_forcing_vars + 1)
             * (num_past_forcing_steps + num_future_forcing_steps + 1)
         )
 
@@ -144,7 +144,7 @@ def __init__(
             self.boundary_dim = (
                 boundary_static_dim
                 # Temporal Embedding counts as one additional forcing_feature
-                + 2*(num_boundary_forcing_vars)
+                + (num_boundary_forcing_vars + 1)
                 * (num_past_boundary_steps + num_future_boundary_steps + 1)
             )