ImportError: cannot import name 'colo_state_dict' from 'colossalai.utils.model.colo_init_context'

[fuhengwu2021]

🐛 Describe the bug

I am trying example colo_vit but got this error

Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
ImportError: cannot import name 'colo_state_dict' from 'colossalai.utils.model.colo_init_context' (/home/wfh/.local/lib/python3.8/site-packages/colossalai/utils/model/colo_init_context.py)

This line seems having issue:

https://github.com/hpcaitech/ColossalAI-Examples/blob/main/image/vision_transformer/colo_vit/train.py#L11

Environment

>>> colossalai.__version__
'0.1.9'

[fuhengwu2021]

I installed colossalai with pip install colossalai==0.1.10+torch1.12cu11.3 -f https://release.colossalai.org

not sure why I got version 0.1.9 as displayed above?