def test_loading_albert():
    """
    Albert is a special case because it includes a Linear layer in the encoder
    that maps the embeddings to the encoder hidden size, but we include this linear
    layer within our embedding layer.
    """
    transformer_embedding = TransformerEmbeddings.from_pretrained_module(
        "albert-base-v2",
    )
    albert = AutoModel.from_pretrained("albert-base-v2")
    assert_allclose(
        transformer_embedding.embeddings.word_embeddings.weight.data,
        albert.embeddings.word_embeddings.weight.data,
    )
    assert_allclose(
        transformer_embedding.linear_transform.weight.data,
        albert.encoder.embedding_hidden_mapping_in.weight.data,
    )
Albert is a special case because it includes a Linear layer in the encoder
that maps the embeddings to the encoder hidden size, but we include this linear
layer within our embedding layer.
tests/modules/transformer/transformer_embeddings_test.py