import torch


def load_optim(params, model, checkpoint=None):
    """
    load optimizer according to specified parameters.
    If checkpoint is specified, restore optim parameters stored checkpoint
    :param params: dict containing parameters of the model
    :param checkpoint : dictionnary. Must contain key 'optim', containing optim.state_dict()
    :param model_dir: if specified, load checkpoint_last.pt state dictionnary
    :return: optimizer
    """
    if params.optim in ('adam', 'Adam', 'ADAM'):
        optim = torch.optim.Adam(model.parameters(), lr=params.learning_rate)
    else:
        # use same params as in DnCNN
        optim = torch.optim.SGD(model.parameters(), lr=params.learning_rate, momentum=0.9)
    if checkpoint is not None:
        optim.load_state_dict(checkpoint['optim'])
    return optim