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