def count_layers(model, trainable_only=True):
    """Count the number of layers in a Keras model.

    Parameters
    ----------

    model : keras.engine.training.Model
        A Keras model.
    trainable_only : bool, optional
        Whether to count only trainable layers.

    Returns
    -------
    int
        Number of layers in the model.
    """
    if trainable_only:
        return sum([c.trainable for c in model.layers])
    return len(model.layers)