AttributeError: Can't pickle local object 'main.<locals>.optimizer_hook

[radhacr]

Seeing the following error while saving the model, while trying to pretrain the mistral-7b model with the pile-of-law dataset

Traceback (most recent call last):
  File "/data/qgalore/Q-GaLore/run_pretrain_mistral.py", line 688, in <module>
    main(args)
  File "/data/qgalore/Q-GaLore/run_pretrain_mistral.py", line 532, in main
    saving_model_weight(model.module, f"{current_model_directory}/pytorch_model.bin")
  File "/data/qgalore/Q-GaLore/q_galore_torch/utils/setup.py", line 38, in saving_model_weight
    torch.save(checkpoint, path)
  File "/data/qgalore/Q-GaLore/venv/lib/python3.10/site-packages/torch/serialization.py", line 628, in save
    _save(obj, opened_zipfile, pickle_module, pickle_protocol, _disable_byteorder_record)
  File "/data/qgalore/Q-GaLore/venv/lib/python3.10/site-packages/torch/serialization.py", line 840, in _save
    pickler.dump(obj)
AttributeError: Can't pickle local object 'main.<locals>.optimizer_hook'
[rank0]: Traceback (most recent call last):
[rank0]:   File "/data/qgalore/Q-GaLore/run_pretrain_mistral.py", line 688, in <module>
[rank0]:     main(args)
[rank0]:   File "/data/qgalore/Q-GaLore/run_pretrain_mistral.py", line 532, in main
[rank0]:     saving_model_weight(model.module, f"{current_model_directory}/pytorch_model.bin")
[rank0]:   File "/data/qgalore/Q-GaLore/q_galore_torch/utils/setup.py", line 38, in saving_model_weight
[rank0]:     torch.save(checkpoint, path)
[rank0]:   File "/data/qgalore/Q-GaLore/venv/lib/python3.10/site-packages/torch/serialization.py", line 628, in save
[rank0]:     _save(obj, opened_zipfile, pickle_module, pickle_protocol, _disable_byteorder_record)
[rank0]:   File "/data/qgalore/Q-GaLore/venv/lib/python3.10/site-packages/torch/serialization.py", line 840, in _save
[rank0]:     pickler.dump(obj)
[rank0]: AttributeError: Can't pickle local object 'main.<locals>.optimizer_hook'