diff --git a/laborious/utils/repository/model_repository.py b/laborious/utils/repository/model_repository.py index 33611bb..11944cc 100644 --- a/laborious/utils/repository/model_repository.py +++ b/laborious/utils/repository/model_repository.py @@ -27,7 +27,7 @@ import gzip import pickle from numpy import ndarray from typing import Any - +import gc from sientia_do.temporal.constants import DATETIME_FORMAT_WITH_TZ ARTIFACTS_PATH = "./tmp/artifacts" @@ -613,8 +613,12 @@ class MLFlowRepository(): prediction = model.predict(data) if retention == 0: + self.logger.info( + f"Deleting model {model_name} from memory") del model + gc.collect() + return prediction def get_cached_predict(self, model_name: str, data: pd.DataFrame, retention: int, @@ -638,8 +642,12 @@ class MLFlowRepository(): prediction = model.predict(data) if retention == 0: + self.logger.info( + f"Deleting model {model_name} from memory") del model + gc.collect() + return prediction """