def convert_to_amp(model: nn.Module, optimizer: Optimizer, criterion: _Loss, mode: AMP_TYPE, amp_config: Config=None):
    'A helper function to wrap training components with Torch AMP modules.\n\n    Args:\n        param model (:class:`torch.nn.Module`): your model object.\n        optimizer (:class:`torch.optim.Optimizer`): your optimizer object.\n        criterion (:class:`torch.nn.modules.loss._Loss`): your loss function object.\n        mode (:class:`colossalai.amp.AMP_TYPE`): amp mode.\n        amp_config (Union[:class:`colossalai.context.Config`, dict]): configuration for different amp modes.\n\n    Returns:\n        A tuple (model, optimizer, criterion).\n\n    Note:\n        ``amp_config`` may vary from different mode you choose. You should check the corresponding amp mode\n        for more details about ``amp_config``.\n        For ``apex_amp``, please check\n        `apex_amp config <https://nvidia.github.io/apex/amp.html?highlight=apex%20amp>`_.\n        For ``naive_amp``, please check\n        `naive_amp config <https://github.com/hpcaitech/ColossalAI/blob/main/colossalai/amp/naive_amp/_fp16_optimizer.py#L42>`_.\n        For ``torch_amp``, please check\n        `torch_amp config <https://github.com/pytorch/pytorch/blob/master/torch/cuda/amp/grad_scaler.py#L97>`_.\n    '
    assert isinstance(mode, AMP_TYPE), f'expected the argument mode be AMP_TYPE, but got {type(mode)}'
    if (amp_config is None):
        amp_config = Config()
    if (mode == AMP_TYPE.TORCH):
        (model, optimizer, criterion) = convert_to_torch_amp(model, optimizer, criterion, amp_config)
    elif (mode == AMP_TYPE.APEX):
        (model, optimizer) = convert_to_apex_amp(model, optimizer, amp_config)
    elif (mode == AMP_TYPE.NAIVE):
        (model, optimizer) = convert_to_naive_amp(model, optimizer, amp_config)
    return (model, optimizer, criterion)
A helper function to wrap training components with Torch AMP modules.

Args:
    param model (:class:`torch.nn.Module`): your model object.
    optimizer (:class:`torch.optim.Optimizer`): your optimizer object.
    criterion (:class:`torch.nn.modules.loss._Loss`): your loss function object.
    mode (:class:`colossalai.amp.AMP_TYPE`): amp mode.
    amp_config (Union[:class:`colossalai.context.Config`, dict]): configuration for different amp modes.

Returns:
    A tuple (model, optimizer, criterion).

Note:
    ``amp_config`` may vary from different mode you choose. You should check the corresponding amp mode
    for more details about ``amp_config``.
    For ``apex_amp``, please check
    `apex_amp config <https://nvidia.github.io/apex/amp.html?highlight=apex%20amp>`_.
    For ``naive_amp``, please check
    `naive_amp config <https://github.com/hpcaitech/ColossalAI/blob/main/colossalai/amp/naive_amp/_fp16_optimizer.py#L42>`_.
    For ``torch_amp``, please check
    `torch_amp config <https://github.com/pytorch/pytorch/blob/master/torch/cuda/amp/grad_scaler.py#L97>`_.
def convert_to_amp(model: nn.Module, optimizer: Optimizer, criterion: _Loss, mode: AMP_TYPE, amp_config: Config=None):
    'A helper function to wrap training components with Torch AMP modules.\n\n    Args:\n        param model (:class:`torch.nn.Module`): your model object.\n        optimizer (:class:`torch.optim.Optimizer`): your optimizer object.\n        criterion (:class:`torch.nn.modules.loss._Loss`): your loss function object.\n        mode (:class:`colossalai.amp.AMP_TYPE`): amp mode.\n        amp_config (Union[:class:`colossalai.context.Config`, dict]): configuration for different amp modes.\n\n    Returns:\n        A tuple (model, optimizer, criterion).\n\n    Note:\n        ``amp_config`` may vary from different mode you choose. You should check the corresponding amp mode\n        for more details about ``amp_config``.\n        For ``apex_amp``, please check\n        `apex_amp config <https://nvidia.github.io/apex/amp.html?highlight=apex%20amp>`_.\n        For ``naive_amp``, please check\n        `naive_amp config <https://github.com/hpcaitech/ColossalAI/blob/main/colossalai/amp/naive_amp/_fp16_optimizer.py#L42>`_.\n        For ``torch_amp``, please check\n        `torch_amp config <https://github.com/pytorch/pytorch/blob/master/torch/cuda/amp/grad_scaler.py#L97>`_.\n    '
    assert isinstance(mode, AMP_TYPE), f'expected the argument mode be AMP_TYPE, but got {type(mode)}'
    if (amp_config is None):
        amp_config = Config()
    if (mode == AMP_TYPE.TORCH):
        (model, optimizer, criterion) = convert_to_torch_amp(model, optimizer, criterion, amp_config)
    elif (mode == AMP_TYPE.APEX):
        (model, optimizer) = convert_to_apex_amp(model, optimizer, amp_config)
    elif (mode == AMP_TYPE.NAIVE):
        (model, optimizer) = convert_to_naive_amp(model, optimizer, amp_config)
    return (model, optimizer, criterion)