def count_parameters(model):
    """Get the number of trainable parameters in the model.

    Parameters
    ----------
    model : nn.Module
        The model

    Returns
    -------
    int
        Number of trainable parameters in the model
    """
    return sum(p.numel() for p in model.parameters() if p.requires_grad)