Skip to content

make wrapped model picklable #1

@arsenyinfo

Description

@arsenyinfo
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'

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions