error when save model llama3.1-70b-instruct

[Daisy-a-p]

I launch pruning this way:

python llama3.py --pruning_ratio 0.25 \
                 --device cuda --base_model meta-llama/Llama-3.1-8B-Instruct \
                 --block_wise --block_mlp_layer_start 4 --block_mlp_layer_end 30 \
                 --block_attention_layer_start 4 --block_attention_layer_end 30 \
                 --save_ckpt_log_name llama3_instruct_prune \
                 --pruner_type l1 \
                 --max_seq_len 2048 \
                 --save_model

When I try to save model llama3.1-70b-instruct, which was pruned on several GPUs, by torch.save(model) I get the following error:

Traceback (most recent call last):
  File "/home/jovyan/voronova/LLM-Pruner/llama3.py", line 366, in <module>
    main(args)
  File "/home/jovyan/voronova/LLM-Pruner/llama3.py", line 277, in main
    torch.save({
  File "/home/jovyan/.mlspace/envs/prune_llm_layers/lib/python3.11/site-packages/torch/serialization.py", line 850, in save
    _save(
  File "/home/jovyan/.mlspace/envs/prune_llm_layers/lib/python3.11/site-packages/torch/serialization.py", line 1088, in _save
    pickler.dump(obj)
AttributeError: Can't pickle local object 'add_hook_to_module.<locals>.new_forward'

How to fix that?