import torch


def load_model_optim(checkpoint_path, model=None, optim=None):
    """Load trained weight.

    Args:
        checkpoint_path (str)
        model (nn.Module)
        optim (torch.optim)

    Returns:
        model
        optim

    """

    ckpt = torch.load(checkpoint_path)
    if model is not None:
        model.load_state_dict(ckpt['model'])
    if optim is not None:
        optim.load_state_dict(ckpt['opt'])
    return model, optim