I’m trying to use joblib inside an airflow dag, when loading a model file I run joblib.load(model_name) but I get an error

  File "/usr/local/lib/python3.9/pickle.py", line 331, in _getattribute
    raise AttributeError("Can't get attribute {!r} on {!r}"
AttributeError: Can't get attribute 'MyCustomModel' on <module '__main__' from '/home/***/.local/bin/***'>

My custom model is imported inside the dag, and I also tried importing inside the task function, but I always get this error that the joblib library can’t find the custom class I need to unpickle the model. I assume this is because the airflow runtime isn’t handling the imports or python namespaces as I expect. Have others tried to load pickles or joblib models in a python dag? Or is there some way to force the runtime to import this custom class?