def METHOD_NAME(path: str,
                    epoch: int,
                    model: torch.nn.Module,
                    optimizer: Optional[ColossalaiOptimizer] = None,
                    lr_scheduler: torch.optim.lr_scheduler._LRScheduler = None,
                    torch_load_kwargs: Optional[Dict] = None,
                    load_state_dict_kwargs: Optional[Dict] = None):
    """load_checkpoint 
    load a model, whose parameters are `ColoTensor`s.
    Args:
        path (str): directory to save the checkpoint files.
        epoch (int): the number of epoch
        model (torch.nn.Module): a torch module initialized by ColoInitContext
        optimizer (ColossalaiOptimizer, optional): optimizers. Defaults to None.
        lr_scheduler (torch.optim.lr_scheduler._LRScheduler, optional): lr schedule. Defaults to None.
        torch_load_kwargs: (dict, optional): The kwargs of torch.load inside the function
        load_state_dict_kwargs (dict, optional): The kwargs of load_state_dict inside the function
    """
    # initialize the default parameters
    if not torch_load_kwargs:
        torch_load_kwargs = dict()
    if not load_state_dict_kwargs:
        load_state_dict_kwargs = dict()
    rank = dist.get_rank()
    mapping = dict()
    for n, p in model.named_parameters():
        if isinstance(p, ColoTensor):
            mapping[n] = p.dist_spec
            gather_tensor(p)
    if rank == 0:
        load_state = torch.load(path + '/epoch_{}_model.pth'.format(epoch), **torch_load_kwargs)
        model.load_state_dict(load_state['model'], **load_state_dict_kwargs)
    dist.barrier()
    # scatter loaded parameters
    for n, p in model.named_parameters():
        if isinstance(p, ColoTensor):
            scatter_tensor(p, mapping[n])
            if rank == 0:
                assert hasattr(p, 'save_ready')
                delattr(p, 'save_ready')
    del mapping
    if optimizer is not None:
        mapping = dict()
        for k, v in optimizer.state_dict()['state'].items():
            for n, t in v.items():
                if isinstance(t, ColoTensor):
                    mapping[(k, n)] = t.dist_spec
                    gather_tensor(t)
        if rank == 0:
            colo_checkpoint = torch.load(path + '/epoch_{}_optim.pth'.format(epoch), **torch_load_kwargs)
            optimizer.load_state_dict(colo_checkpoint['optim'], **load_state_dict_kwargs)
        dist.barrier()
        for k, v in optimizer.state_dict()['state'].items():
            for n, t in v.items():
                if isinstance(t, ColoTensor):
                    scatter_tensor(t, mapping[(k, n)])
        del mapping