def find_embedding_layer(model: torch.nn.Module) -> torch.nn.Module:
    """
    Takes a model (typically an AllenNLP `Model`, but this works for any `torch.nn.Module`) and
    makes a best guess about which module is the embedding layer.  For typical AllenNLP models,
    this often is the `TextFieldEmbedder`, but if you're using a pre-trained contextualizer, we
    really want layer 0 of that contextualizer, not the output.  So there are a bunch of hacks in
    here for specific pre-trained contextualizers.
    """
    # We'll look for a few special cases in a first pass, then fall back to just finding a
    # TextFieldEmbedder in a second pass if we didn't find a special case.
    from transformers.models.gpt2.modeling_gpt2 import GPT2Model
    from transformers.models.bert.modeling_bert import BertEmbeddings
    from transformers.models.albert.modeling_albert import AlbertEmbeddings
    from transformers.models.roberta.modeling_roberta import RobertaEmbeddings
    from allennlp.modules.text_field_embedders.text_field_embedder import TextFieldEmbedder
    from allennlp.modules.text_field_embedders.basic_text_field_embedder import (
        BasicTextFieldEmbedder,
    )
    from allennlp.modules.token_embedders.embedding import Embedding

    for module in model.modules():
        if isinstance(module, BertEmbeddings):
            return module.word_embeddings
        if isinstance(module, RobertaEmbeddings):
            return module.word_embeddings
        if isinstance(module, AlbertEmbeddings):
            return module.word_embeddings
        if isinstance(module, GPT2Model):
            return module.wte

    for module in model.modules():
        if isinstance(module, TextFieldEmbedder):

            if isinstance(module, BasicTextFieldEmbedder):
                # We'll have a check for single Embedding cases, because we can be more efficient
                # in cases like this.  If this check fails, then for something like hotflip we need
                # to actually run the text field embedder and construct a vector for each token.
                if len(module._token_embedders) == 1:
                    embedder = list(module._token_embedders.values())[0]
                    if isinstance(embedder, Embedding):
                        if embedder._projection is None:
                            # If there's a projection inside the Embedding, then we need to return
                            # the whole TextFieldEmbedder, because there's more computation that
                            # needs to be done than just multiply by an embedding matrix.
                            return embedder
            return module
    raise RuntimeError("No embedding module found!")
Takes a model (typically an AllenNLP `Model`, but this works for any `torch.nn.Module`) and
makes a best guess about which module is the embedding layer.  For typical AllenNLP models,
this often is the `TextFieldEmbedder`, but if you're using a pre-trained contextualizer, we
really want layer 0 of that contextualizer, not the output.  So there are a bunch of hacks in
here for specific pre-trained contextualizers.
allennlp/nn/util.py