[BUG]: AttributeError: 'ColossalAIStrategy' object has no attribute 'model'

[grib0ed0v]

🐛 Describe the bug

When you run training with colossalai_zero2 or colossalai_gemini strategy, after training you cannot save the model due to:

Traceback (most recent call last):
  File "train_prompts.py", line 132, in <module>
    main(args)
  File "train_prompts.py", line 106, in main
    strategy.save_model(actor, args.save_path, only_rank0=True)
  File "/home/agruzdev/.virtualenvs/colossalai/lib/python3.8/site-packages/chatgpt-1.0.0-py3.8.egg/chatgpt/trainer/strategies/colossalai.py", line 176, in save_model
AttributeError: 'ColossalAIStrategy' object has no attribute 'model'

Environment

Python 3.8
PyTorch 1.13
CUDA 11.4