diff --git a/mlair/data_handler/iterator.py b/mlair/data_handler/iterator.py
index e353f84d85a0871b00964899efb2a79bf555aefc..3fc25a90f861c65d38aa6b7019095210035d4c2d 100644
--- a/mlair/data_handler/iterator.py
+++ b/mlair/data_handler/iterator.py
@@ -144,8 +144,8 @@ class KerasIterator(keras.utils.Sequence):
         mod_rank = self._get_model_rank()
         for data in self._collection:
             logging.debug(f"prepare batches for {str(data)}")
-            X = data.get_X(upsampling=self.upsampling)
-            Y = [data.get_Y(upsampling=self.upsampling)[0] for _ in range(mod_rank)]
+            X, _Y = data.get_data(upsampling=self.upsampling)
+            Y = [_Y[0] for _ in range(mod_rank)]
             if self.upsampling:
                 X, Y = self._permute_data(X, Y)
             if remaining is not None: