def total_parameter_size(model):
    """look up state dict of the model and returns the total number of
    parameters.

    Parameters:
        model (nn.Module): pytorch model.
    Returns:
        the number of parameters including custom weight matrices.
    """
    return sum([w.nelement() for k, w in model.state_dict().items()])