diff --git a/mlair/run_modules/training.py b/mlair/run_modules/training.py
index 113765e0d295bb0b1d756cd1cefba85093b20089..3878e79cb66365ba186ee736f3a4927c076d2dee 100644
--- a/mlair/run_modules/training.py
+++ b/mlair/run_modules/training.py
@@ -145,7 +145,7 @@ class Training(RunEnvironment):
             logging.info("Found locally stored model and checkpoints. Training is resumed from the last checkpoint.")
             self.callbacks.load_callbacks()
             self.callbacks.update_checkpoint()
-            self.model = keras.models.load_model(checkpoint.filepath)
+            self.model = keras.models.load_model(checkpoint.filepath, self.model.custom_objects)
             hist: History = self.callbacks.get_callback_by_name("hist")
             initial_epoch = max(hist.epoch) + 1
             _ = self.model.fit_generator(generator=self.train_set,