src/transformers/models/megatron_bert/modeling_megatron_bert.py/MegatronBertAttention/__init__
class MegatronBertAttention:    def __init__(self, config):
        super().__init__()
        self.ln = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
        self.self = MegatronBertSelfAttention(config)
        self.output = MegatronBertSelfOutput(config)
        self.pruned_heads = set()