def get_params(model, only_trainable=False):
    """
    Get the parameter number of the model.
    If only_trainable is true, only trainable parameters will be counted.
    """

    if not only_trainable:
        return sum(p.numel() for p in model.parameters())
    else:
        return sum(p.numel() for p in model.parameters() if p.requires_grad)