src/transformers/models/megatron_bert/modeling_megatron_bert.py/MegatronBertPredictionHeadTransform/forward
class MegatronBertPredictionHeadTransform:    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
        hidden_states = self.dense(hidden_states)
        hidden_states = self.transform_act_fn(hidden_states)
        hidden_states = self.LayerNorm(hidden_states)
        return hidden_states