def test_labram_returns(default_labram_params, use_mean_pooling):
    """
    Testing if the model is returning the correct shapes for the different
    return options.

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

    """
    labram_base = Labram(
        n_layers=12,
        att_num_heads=12,
        use_mean_pooling=use_mean_pooling,
        **default_labram_params,
    )
    # Defining a random data
    X = torch.rand(1, default_labram_params["n_chans"],
                   default_labram_params["n_times"])

    with torch.no_grad():
        out = labram_base(X, return_all_tokens=False,
                          return_patch_tokens=False)

        assert out.shape == torch.Size([1, default_labram_params["n_outputs"]])

        out_patches = labram_base(X, return_all_tokens=False,
                                  return_patch_tokens=True)

        assert out_patches.shape == torch.Size(
            [1, 320, default_labram_params["n_outputs"]]
        )

        out_all_tokens = labram_base(X, return_all_tokens=True,
                                     return_patch_tokens=False)
        assert out_all_tokens.shape == torch.Size(
            [1, 321, default_labram_params["n_outputs"]]
        )

    Testing if the model is returning the correct shapes for the different
    return options.

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

    
test/unit_tests/models/test_models.py