Saving and restoring checkpoint hangs in Jupyter Notebook. about orbax HOT 4 CLOSED

I have a suspicion it may be due to a bad interaction between different dependencies. Could you try to strip out the extra stuff unrelated to checkpointing, without removing any imports and see if that works?

If it still hangs, you can start removing some imports, and the issue should resolve (at worst by the point where you're only importing Orbax).

Also, a few nits:
handlers = {'state': orbax.checkpoint.PyTreeCheckpointer()}: this should be checkpointers. Handler is a different concept in Orbax and will rapidly lead to bugs if you confuse the two.
orbax.checkpoint.Checkpointer(orbax.checkpoint.PyTreeCheckpointHandler()) is equivalent to orbax.checkpoint.PyTreeCheckpointer()

from orbax.

Thank you for the suggestions! It turns out

import matplotlib.pyplot as plt
import matplotlib
matplotlib.use('QtAgg')
%matplotlib qt

in my notebook were the lines of code causing the issue. Is there a way to solve this issue other than not importing?

from orbax.

I'm afraid your guess is as good as mine, since I'm not familiar with this particular backend, and there appears to be no problem with matplotlib without overriding the backend. Is it possible to skip that part?

from orbax.

Yep, matplotlib by itself is not the issue. The qt backend is more of a convenience, not at all necessary. I will disable it for now. Thanks for the help!

from orbax.