Can't load pickle with custom preprocessing classes

Question

I have two folders:

• notebooks (where I've trained the model)
• src/main.py

In main.py I'm trying to load the model using joblib but I get an error that the custom preprocessing classes can't be found. Even though, I've imported it.

import logging
from logging import config
from fastapi import Depends, FastAPI, HTTPException, status

from domain.services import PredictionService
from usecase.item import ItemReadModel
from utils.preprocess import PriceImputer

import joblib

# config.fileConfig("logging.conf", disable_existing_loggers=False)
logger = logging.getLogger(__name__)

app = FastAPI()

model = joblib.load("models/logistic_regression_model.joblib")

I get this error:

AttributeError: Can't get attribute 'PriceImputer' on <module '__mp_main__' from '...'>

But as you may see, I've already imported it.

[EDIT] I just needed to load the model under the if __name__ == '__main__':.

if __name__ == "__main__":
    model = joblib.load("models/logistic_regression_model.joblib")
    prediction_service = PredictionService(model)
    uvicorn.run(app, host='0.0.0.0', port=80)