implement index in field

ecmwf-lab · Sep 14, 2024 · 3068f2f · 3068f2f
1 parent 5eab711
commit 3068f2f
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/src/ai_models/inputs/transform.py b/src/ai_models/inputs/transform.py
@@ -21,6 +21,8 @@ def to_numpy(self, flatten=False, dtype=None, index=None):
             data = data.astype(dtype)
         if flatten:
             data = data.flatten()
+        if index is not None:
+            data = data[index]
         return data
 
     def __getattr__(self, name):