Forgot to push this change

MeteoSwiss · Oct 18, 2024 · 6d4e127 · 6d4e127
1 parent 041d64e
commit 6d4e127
Showing 1 changed file with 6 additions and 2 deletions.
diff --git a/src/anemoi/training/data/dataset.py b/src/anemoi/training/data/dataset.py
@@ -239,8 +239,12 @@ def __iter__(self) -> torch.Tensor:
             x = self.data[start : end : self.timeincrement]
             if self.spatial_index is not None:
                 x = x[..., self.spatial_index]
-            x = rearrange(x, "dates variables ensemble gridpoints -> dates ensemble gridpoints variables")
-            self.ensemble_dim = 1
+            if self.label == "predict":
+                x = rearrange(x, "dates variables ensemble gridpoints -> ensemble dates gridpoints variables")
+                self.ensemble_dim = 0
+            else:
+                x = rearrange(x, "dates variables ensemble gridpoints -> dates ensemble gridpoints variables")
+                self.ensemble_dim = 1
 
             yield torch.from_numpy(x)