Fix boundary masking bug for static features

mllam · Dec 2, 2024 · 29cc25b · 29cc25b
1 parent b35072d
commit 29cc25b
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
@@ -70,12 +70,12 @@ def __init__(
         static_features_torch = torch.tensor(arr_static, dtype=torch.float32)
         self.register_buffer(
             "grid_static_features",
-            static_features_torch[self.boundary_mask[:, 0].to(torch.bool)],
+            static_features_torch[self.interior_mask[:, 0].to(torch.bool)],
             persistent=False,
         )
         self.register_buffer(
             "boundary_static_features",
-            static_features_torch[self.interior_mask[:, 0].to(torch.bool)],
+            static_features_torch[self.boundary_mask[:, 0].to(torch.bool)],
             persistent=False,
         )