def test_model_trainable_parameters_labram(default_labram_params):
    """
    Test the number of trainable parameters in Labram model based on the
    paper values.

    Parameters
    ----------
    default_labram_params: dict with default parameters for Labram model

    """
    labram_base = Labram(n_layers=12, att_num_heads=12,
                         **default_labram_params)

    labram_base_parameters = labram_base.get_torchinfo_statistics().trainable_params

    # We added some parameters layers in the segmentation step to match the
    # braindecode convention.
    assert np.round(labram_base_parameters / 1e6, 1) == 5.8
    # ~ 5.8 M matching the paper

    labram_large = Labram(
        n_layers=24,
        att_num_heads=16,
        out_channels=16,
        emb_size=400,
        **default_labram_params,
    )
    labram_large_parameters = labram_large.get_torchinfo_statistics().trainable_params

    assert np.round(labram_large_parameters / 1e6, 0) == 46
    # ~ 46 M matching the paper

    labram_huge = Labram(
        n_layers=48,
        att_num_heads=16,
        out_channels=32,
        emb_size=800,
        **default_labram_params,
    )

    labram_huge_parameters = labram_huge.get_torchinfo_statistics().trainable_params
    # 369M matching the paper
    assert np.round(labram_huge_parameters / 1e6, 0) == 369

    assert labram_base.get_num_layers() == 12
    assert labram_large.get_num_layers() == 24
    assert labram_huge.get_num_layers() == 48

    Test the number of trainable parameters in Labram model based on the
    paper values.

    Parameters
    ----------
    default_labram_params: dict with default parameters for Labram model

    
test/unit_tests/models/test_models.py