diff --git a/src/anemoi/datasets/data/grids.py b/src/anemoi/datasets/data/grids.py
index 15915f1..fd57413 100644
--- a/src/anemoi/datasets/data/grids.py
+++ b/src/anemoi/datasets/data/grids.py
@@ -324,7 +324,8 @@ def shape(self):
         """
         shapes = [np.sum(mask) for mask in self.masks]
         global_shape = np.sum(self.global_mask)
-        return tuple(self.lams[0].shape[:-1] + (sum(shapes) + global_shape,))
+        total_shape = sum(shapes) + global_shape
+        return tuple(self.lams[0].shape[:-1] + (int(total_shape),))
 
     def check_same_resolution(self, d1, d2):
         # Turned off because we are combining different resolutions