diff --git a/mlair/run_modules/post_processing.py b/mlair/run_modules/post_processing.py
index 71c49433f34949b538423312c10152d03312165f..5e1e585e2114fe29cf01bb89bbbccb17d7bfa4bf 100644
--- a/mlair/run_modules/post_processing.py
+++ b/mlair/run_modules/post_processing.py
@@ -68,7 +68,7 @@ class PostProcessing(RunEnvironment):
     def __init__(self):
         """Initialise and run post-processing."""
         super().__init__()
-        self.model: keras.Model = self._load_model()
+        self.model: AbstractModelClass = self._load_model()
         self.model_name = self.data_store.get("model_name", "model").rsplit("/", 1)[1].split(".", 1)[0]
         self.ols_model = None
         self.batch_size: int = self.data_store.get_default("batch_size", "model", 64)
@@ -432,8 +432,8 @@ class PostProcessing(RunEnvironment):
         except NameNotFoundInDataStore:
             logging.info("No model was saved in data store. Try to load model from experiment path.")
             model_name = self.data_store.get("model_name", "model")
-            model_class: AbstractModelClass = self.data_store.get("model", "model")
-            model = keras.models.load_model(model_name, custom_objects=model_class.custom_objects)
+            model: AbstractModelClass = self.data_store.get("model", "model")
+            model.load_model(model_name)
         return model
 
     # noinspection PyBroadException