works standalone

MeteoSwiss · May 2, 2024 · 0f5217a · 0f5217a
1 parent df9b595
commit 0f5217a
Showing 1 changed file with 7 additions and 2 deletions.
diff --git a/neural_lam/models/ar_model.py b/neural_lam/models/ar_model.py
@@ -879,7 +879,12 @@ def save_pred_as_grib(self, file_path, value_dir_path):
             for variable, grib_code in constants.GRIB_NAME.items():
                 # here find the key of the cariable in constants.is_3D
                 #  and if == 7, assign a cut of 7 on the reshape. Else 1
-                shape_val = 13 if constants.IS_3D[variable] else 1
+                if constants.IS_3D[variable]:
+                    shape_val = 13 
+                    vertical = constants.VERTICAL_LEVELS
+                else:
+                    shape_val = 1
+                    vertical = 1
                 # Find the value range to sample
                 value_range = indices[variable]
 
@@ -891,7 +896,7 @@ def save_pred_as_grib(self, file_path, value_dir_path):
                 original_data = earthkit.data.from_source("file", sample_file)
 
                 subset = original_data.sel(
-                    shortName=grib_code, level=constants.VERTICAL_LEVELS
+                    shortName=grib_code, level=vertical
                 )
                 md = subset.metadata()