Importing Bagua modules do no work

[silverCore97]

Describe the bug
While intending to run my Bagua code on google colab, the error appeared already with this import from bagua.torch_api.bucket import BaguaBucket.
The error itself is "PydanticUserError: Please use typing_extensions.TypedDict instead of typing.TypedDict on Python < 3.12." But in bagua.bagua_define.py we have

if sys.version_info >= (3, 9):
    from typing import TypedDict  # pytype: disable=not-supported-yet
else:
    from typing_extensions import TypedDict  # pytype: disable=not-supported-yet
from pydantic import BaseModel

Environment

• Your operating system and version: "Ubuntu 22.04.3 LTS"
• Your python version: 3.10.12
• Your PyTorch version: not relevant
• How did you install python (e.g. apt or pyenv)? Did you use a virtualenv?: Comes pre-installed in Google Colab; No
• Have you tried using latest bagua master (python3 -m pip install --pre bagua)?: Yes, does not work either

Reproducing

Please provide a minimal working example. This means the runnable code.

!pip install bagua-cuda116
from bagua.torch_api.bucket import BaguaBucket

Please also write what exact commands are required to reproduce your results.

Additional context
It''s on Google Colab.