hookandlook icon indicating copy to clipboard operation
hookandlook copied to clipboard

make wrapped model picklable

Open arsenyinfo opened this issue 3 years ago • 0 comments

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'

arsenyinfo avatar Feb 18 '22 16:02 arsenyinfo