src/transformers/models/esm/modeling_tf_esm.py/TFEsmLMHead/call
class TFEsmLMHead:    def call(self, features):
        x = self.dense(features)
        x = gelu(x)
        x = self.layer_norm(x)

        # project back to size of vocabulary with bias
        x = self.decoder(x)
        x = x + self.bias
        return x