src/transformers/models/megatron_bert/modeling_megatron_bert.py/MegatronBertLayer/feed_forward_chunk
class MegatronBertLayer:    def feed_forward_chunk(self, attention_output):
        ln_output = self.ln(attention_output)
        intermediate_output = self.intermediate(ln_output)
        layer_output = self.output(intermediate_output, attention_output)
        return layer_output