update predict to ensure it has the correct attributes on pred_pr

henryaddison · Dec 14, 2023 · a32414d · a32414d
1 parent f48d2b5
commit a32414d
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/bin/predict.py b/bin/predict.py
@@ -165,6 +165,7 @@ def np_samples_to_xr(np_samples, target_transform, coords, cf_data_vars):
         xr.Dataset(data_vars=data_vars, coords=coords, attrs={})
     )
     samples_ds = samples_ds.rename({"target_pr": "pred_pr"})
+    samples_ds["pred_pr"] = samples_ds["pred_pr"].assign_attrs(pred_pr_attrs)
     return samples_ds