diff --git a/mlair/data_handler/default_data_handler.py b/mlair/data_handler/default_data_handler.py
index 1bcc1491652c630a0059e00ee92406ca3d4df495..300e0435c4e8441e299675319e2c72604ebb3200 100644
--- a/mlair/data_handler/default_data_handler.py
+++ b/mlair/data_handler/default_data_handler.py
@@ -125,8 +125,9 @@ class DefaultDataHandler(AbstractDataHandler):
 
     def get_data(self, upsampling=False, as_numpy=True):
         self._load()
-        X = self.get_X(upsampling, as_numpy)
-        Y = self.get_Y(upsampling, as_numpy)
+        as_numpy_X, as_numpy_Y = as_numpy if isinstance(as_numpy, tuple) else (as_numpy, as_numpy)
+        X = self.get_X(upsampling, as_numpy_X)
+        Y = self.get_Y(upsampling, as_numpy_Y)
         self._reset_data()
         return X, Y
 
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:
diff --git a/mlair/run_modules/post_processing.py b/mlair/run_modules/post_processing.py
index 07ef1ce46b7951c812fbd17cc06916bbd8cb9caf..00d82f3c6f48c3560e31d62b5bed4ddbd2bc49be 100644
--- a/mlair/run_modules/post_processing.py
+++ b/mlair/run_modules/post_processing.py
@@ -10,7 +10,6 @@ import sys
 import traceback
 from typing import Dict, Tuple, Union, List, Callable
 
-import tensorflow.keras as keras
 import numpy as np
 import pandas as pd
 import xarray as xr
@@ -695,6 +694,7 @@ class PostProcessing(RunEnvironment):
         logging.info(f"start train_ols_model on train data")
         self.ols_model = OrdinaryLeastSquaredModel(self.train_data)
 
+    @TimeTrackingWrapper
     def make_prediction(self, subset):
         """
         Create predictions for NN, OLS, and persistence and add true observation as reference.
@@ -707,7 +707,7 @@ class PostProcessing(RunEnvironment):
         logging.info(f"start make_prediction for {subset_type}")
         time_dimension = self.data_store.get("time_dim")
         window_dim = self.data_store.get("window_dim")
-        subset_type = subset.name
+
         for i, data in enumerate(subset):
             input_data = data.get_X()
             target_data = data.get_Y(as_numpy=False)
diff --git a/test/test_data_handler/test_iterator.py b/test/test_data_handler/test_iterator.py
index e47d725a4fd78fec98e81a6de9c18869e7b47637..bb8ecb5d216519b3662a5baa4d463780b4c29d8c 100644
--- a/test/test_data_handler/test_iterator.py
+++ b/test/test_data_handler/test_iterator.py
@@ -106,6 +106,9 @@ class DummyData:
         Y2 = np.random.randint(21, 30, size=(self.number_of_samples, 5, 1))  # samples, window, variables
         return [Y1, Y2]
 
+    def get_data(self, upsampling=False, as_numpy=True):
+        return self.get_X(upsampling, as_numpy), self.get_Y(upsampling, as_numpy)
+
 
 class TestKerasIterator:
 
diff --git a/test/test_run_modules/test_model_setup.py b/test/test_run_modules/test_model_setup.py
index 962287e09aacd3c44961a827c86b331d643ec401..6e8d3ea9ebab40c79b17b2fba386322a630f00e1 100644
--- a/test/test_run_modules/test_model_setup.py
+++ b/test/test_run_modules/test_model_setup.py
@@ -150,3 +150,6 @@ class DummyData:
         Y1 = np.random.randint(0, 10, size=(self.number_of_samples, 5))  # samples, window
         Y2 = np.random.randint(21, 30, size=(self.number_of_samples, 3))  # samples, window
         return [Y1, Y2]
+
+    def get_data(self, upsampling=False, as_numpy=True):
+        return self.get_X(upsampling, as_numpy), self.get_Y(upsampling, as_numpy)