import importlib


def set_optimizer(model, optimizer_name, **kwargs):
    """
    Set the parameter optimizer for the model.

    Reference: https://pytorch.org/docs/stable/optim.html#algorithms
    """

    torch_optim_optimizers = [
        "Adadelta",
        "Adagrad",
        "Adam",
        "AdamW",
        "Adamax",
        "ASGD",
        "RMSprop",
        "Rprop",
        "SGD",
    ]
    if optimizer_name not in torch_optim_optimizers:
        msg = "Unrecognized optimizer: {}, should be one of {}."
        raise NotImplementedError(
            msg.format(optimizer_name, ",".join(torch_optim_optimizers))
        )

    optimizer_cls = getattr(
        importlib.import_module("torch.optim"), optimizer_name
    )
    optimizer = optimizer_cls(model.parameters(), **kwargs)

    return optimizer