ploomber: Partial build fails - TypeError: cannot pickle 'module' object

I see the following error when I do ploomber build -p train, but do not get the error when running ploomber build. My guess it that it relates to a joblib.dump. I tried to add custom un/serializers after seeing the error, but that did not fix it. Sorry, the pipeline is a bit complicated at this point and the error is hard to generate a easy to reproduce example for testing.

Traceback (most recent call last):
  File "/home/hopkir/anaconda3/envs/hydrates/lib/python3.8/site-packages/ploomber/cli/io.py", line 34, in wrapper
    fn(**kwargs)
  File "/home/hopkir/anaconda3/envs/hydrates/lib/python3.8/site-packages/ploomber/cli/build.py", line 49, in main
    report = dag.build_partially(args.partially,
  File "/home/hopkir/anaconda3/envs/hydrates/lib/python3.8/site-packages/ploomber/dag/dag.py", line 667, in build_partially
    dag_copy = deepcopy(self)
  File "/home/hopkir/anaconda3/envs/hydrates/lib/python3.8/copy.py", line 172, in deepcopy
    y = _reconstruct(x, memo, *rv)
  File "/home/hopkir/anaconda3/envs/hydrates/lib/python3.8/copy.py", line 270, in _reconstruct
    state = deepcopy(state, memo)
  File "/home/hopkir/anaconda3/envs/hydrates/lib/python3.8/copy.py", line 146, in deepcopy
    y = copier(x, memo)
  File "/home/hopkir/anaconda3/envs/hydrates/lib/python3.8/copy.py", line 230, in _deepcopy_dict
    y[deepcopy(key, memo)] = deepcopy(value, memo)
  File "/home/hopkir/anaconda3/envs/hydrates/lib/python3.8/copy.py", line 172, in deepcopy
    y = _reconstruct(x, memo, *rv)
  File "/home/hopkir/anaconda3/envs/hydrates/lib/python3.8/copy.py", line 270, in _reconstruct
    state = deepcopy(state, memo)
  File "/home/hopkir/anaconda3/envs/hydrates/lib/python3.8/copy.py", line 146, in deepcopy
    y = copier(x, memo)
  File "/home/hopkir/anaconda3/envs/hydrates/lib/python3.8/copy.py", line 230, in _deepcopy_dict
    y[deepcopy(key, memo)] = deepcopy(value, memo)
  File "/home/hopkir/anaconda3/envs/hydrates/lib/python3.8/copy.py", line 146, in deepcopy
    y = copier(x, memo)
  File "/home/hopkir/anaconda3/envs/hydrates/lib/python3.8/copy.py", line 230, in _deepcopy_dict
    y[deepcopy(key, memo)] = deepcopy(value, memo)
  File "/home/hopkir/anaconda3/envs/hydrates/lib/python3.8/copy.py", line 146, in deepcopy
    y = copier(x, memo)
  File "/home/hopkir/anaconda3/envs/hydrates/lib/python3.8/copy.py", line 230, in _deepcopy_dict
    y[deepcopy(key, memo)] = deepcopy(value, memo)
  File "/home/hopkir/anaconda3/envs/hydrates/lib/python3.8/copy.py", line 172, in deepcopy
    y = _reconstruct(x, memo, *rv)
  File "/home/hopkir/anaconda3/envs/hydrates/lib/python3.8/copy.py", line 270, in _reconstruct
    state = deepcopy(state, memo)
  File "/home/hopkir/anaconda3/envs/hydrates/lib/python3.8/copy.py", line 146, in deepcopy
    y = copier(x, memo)
  File "/home/hopkir/anaconda3/envs/hydrates/lib/python3.8/copy.py", line 230, in _deepcopy_dict
    y[deepcopy(key, memo)] = deepcopy(value, memo)
  File "/home/hopkir/anaconda3/envs/hydrates/lib/python3.8/copy.py", line 172, in deepcopy
    y = _reconstruct(x, memo, *rv)
  File "/home/hopkir/anaconda3/envs/hydrates/lib/python3.8/copy.py", line 270, in _reconstruct
    state = deepcopy(state, memo)
  File "/home/hopkir/anaconda3/envs/hydrates/lib/python3.8/copy.py", line 146, in deepcopy
    y = copier(x, memo)
  File "/home/hopkir/anaconda3/envs/hydrates/lib/python3.8/copy.py", line 230, in _deepcopy_dict
    y[deepcopy(key, memo)] = deepcopy(value, memo)
  File "/home/hopkir/anaconda3/envs/hydrates/lib/python3.8/copy.py", line 172, in deepcopy
    y = _reconstruct(x, memo, *rv)
  File "/home/hopkir/anaconda3/envs/hydrates/lib/python3.8/copy.py", line 270, in _reconstruct
    state = deepcopy(state, memo)
  File "/home/hopkir/anaconda3/envs/hydrates/lib/python3.8/copy.py", line 146, in deepcopy
    y = copier(x, memo)
  File "/home/hopkir/anaconda3/envs/hydrates/lib/python3.8/copy.py", line 230, in _deepcopy_dict
    y[deepcopy(key, memo)] = deepcopy(value, memo)
  File "/home/hopkir/anaconda3/envs/hydrates/lib/python3.8/copy.py", line 172, in deepcopy
    y = _reconstruct(x, memo, *rv)
  File "/home/hopkir/anaconda3/envs/hydrates/lib/python3.8/copy.py", line 270, in _reconstruct
    state = deepcopy(state, memo)
  File "/home/hopkir/anaconda3/envs/hydrates/lib/python3.8/copy.py", line 146, in deepcopy
    y = copier(x, memo)
  File "/home/hopkir/anaconda3/envs/hydrates/lib/python3.8/copy.py", line 230, in _deepcopy_dict
    y[deepcopy(key, memo)] = deepcopy(value, memo)
  File "/home/hopkir/anaconda3/envs/hydrates/lib/python3.8/copy.py", line 172, in deepcopy
    y = _reconstruct(x, memo, *rv)
  File "/home/hopkir/anaconda3/envs/hydrates/lib/python3.8/copy.py", line 270, in _reconstruct
    state = deepcopy(state, memo)
  File "/home/hopkir/anaconda3/envs/hydrates/lib/python3.8/copy.py", line 146, in deepcopy
    y = copier(x, memo)
  File "/home/hopkir/anaconda3/envs/hydrates/lib/python3.8/copy.py", line 230, in _deepcopy_dict
    y[deepcopy(key, memo)] = deepcopy(value, memo)
  File "/home/hopkir/anaconda3/envs/hydrates/lib/python3.8/copy.py", line 161, in deepcopy
    rv = reductor(4)
TypeError: cannot pickle 'module' object