src/transformers/models/esm/modeling_esm.py/EsmLayer/feed_forward_chunk
class EsmLayer:    def feed_forward_chunk(self, attention_output):
        attention_output_ln = self.LayerNorm(attention_output)
        intermediate_output = self.intermediate(attention_output_ln)
        layer_output = self.output(intermediate_output, attention_output)
        return layer_output