src/transformers/models/megatron_bert/modeling_megatron_bert.py/MegatronBertLMPredictionHead/forward
class MegatronBertLMPredictionHead:    def forward(self, hidden_states):
        hidden_states = self.transform(hidden_states)
        hidden_states = self.decoder(hidden_states)
        return hidden_states