hookandlook
hookandlook copied to clipboard
make wrapped model picklable
In [1]: import torch
In [2]: from hookandlook.watch import Wrapper
In [3]: model = Wrapper.wrap_model(torch.nn.Linear(10, 10))
In [4]: torch.save(model, '/tmp/model.pth')
---------------------------------------------------------------------------
AttributeError Traceback (most recent call last)
<ipython-input-4-1952706372cd> in <module>
----> 1 torch.save(model, '/tmp/model.pth')
~/.pyenv/versions/3.7.8/lib/python3.7/site-packages/torch/serialization.py in save(obj, f, pickle_module, pickle_protocol, _use_new_zipfile_serialization)
368
369 with _open_file_like(f, 'wb') as opened_file:
--> 370 _legacy_save(obj, opened_file, pickle_module, pickle_protocol)
371
372
~/.pyenv/versions/3.7.8/lib/python3.7/site-packages/torch/serialization.py in _legacy_save(obj, f, pickle_module, pickle_protocol)
441 pickler = pickle_module.Pickler(f, protocol=pickle_protocol)
442 pickler.persistent_id = persistent_id
--> 443 pickler.dump(obj)
444
445 serialized_storage_keys = sorted(serialized_storages.keys())
AttributeError: Can't pickle local object 'compose_hooks.<locals>.composite_hook'