flax
b7571e44 - Fix cloudpickle serialization of wrapped Module methods.

Commit
4 years ago
Fix cloudpickle serialization of wrapped Module methods. Calling _context inside a closure will make cloudpickle attempt to serialize. It will fail because it cannot serialize ThreadLocals. Factoring out the body of method_wrapper into a seperate method hides the _context for cloudpickle because it can store library methods by name.
Author
Committer
Parents
Loading