AttributeError: Can't pickle local object 'add_hook_to_module.<locals>.new_forward'

[shwu-nyunai]

Traceback (most recent call last):
File "/home/shwu/LABS/FLAP/main.py", line 147, in
main()
File "/home/shwu/LABS/FLAP/main.py", line 141, in main
torch.save(model, f'{args.save_model}/pruned_model.pt')
File "/home/shwu/LABS/FLAP/venv/lib/python3.10/site-packages/torch/serialization.py", line 629, in save
_save(obj, opened_zipfile, pickle_module, pickle_protocol, _disable_byteorder_record)
File "/home/shwu/LABS/FLAP/venv/lib/python3.10/site-packages/torch/serialization.py", line 841, in _save
pickler.dump(obj)
AttributeError: Can't pickle local object 'add_hook_to_module..new_forward'

---

this occurs after pruning on multiple gpus and saving with torch.save

[NEUBuffett]

Hello, I have the same problem, did you solve it?