Pickle error with Jupyter notebook

FYI, an error which took me quite some time to solve.

I wanted to use ray (ray 0.8.4) to tune the hyper-parameters of a torch (1.5.0) nn model (eg: Foo model) in a Jupyter notebook, with all my Python lines of code (Foo model, torch and ray code) being directly inside the Jupyter notebook.

At the end of the first back-propagation, I always had the following error:
File “/home/xxx/.local/lib/python3.6/site-packages/torch/serialization.py”, line 443, in _legacy_save
pickler.dump(obj)
_pickle.PicklingError: Can’t pickle <class ‘main.Foo’>: attribute lookup Foo on main failed

I successfully fixed it by moving out the python code related the nn model in a separate .py file ( foo.py) and importing it in the notebook (from foo import Foo).