src/transformers/models/esm/modeling_tf_esm.py/TFEsmMainLayer/__init__
class TFEsmMainLayer:    def __init__(self, config, add_pooling_layer=True, name=None, **kwargs):
        super().__init__(name=name, **kwargs)

        self.config = config
        self.is_decoder = config.is_decoder

        self.embeddings = TFEsmEmbeddings(config, name="embeddings")
        self.encoder = TFEsmEncoder(config, name="encoder")
        self.pooler = TFEsmPooler(config, name="pooler") if add_pooling_layer else None

        self.contact_head = TFEsmContactPredictionHead(
            in_features=self.config.num_hidden_layers * self.config.num_attention_heads, bias=True, name="contact_head"
        )