Method __init__

src/bert_layers/layers.py:675–678

(self, config: FlexBertConfig)

    """

    def __init__(self, config: FlexBertConfig):
        super().__init__()
        self.layers = nn.ModuleList([get_bert_layer(config, layer_id=i) for i in range(config.num_hidden_layers)])
        self.num_attention_heads = config.num_attention_heads

    def forward(self, hidden_states: torch.Tensor, attention_mask: torch.Tensor, **kwargs) -> torch.Tensor:
        for layer_module in self.layers: