Method __init__

axlearn/common/bert.py:219–226

(self, cfg: Config, *, parent: Optional[Module])

        output: Linear.Config = Linear.default_config()  # Output projection layer.

    def __init__(self, cfg: Config, *, parent: Optional[Module]):
        super().__init__(cfg, parent=parent)
        cfg = self.config
        self._add_child("pooler", cfg.pooler.set(input_dim=cfg.input_dim))
        self._add_child(
            "output",
            cfg.output.set(input_dim=cfg.input_dim, output_dim=cfg.num_classes),
        )

    def forward(self, input_batch: NestedTensor) -> Tensor:
        """Produces prediction scores from BERT output features.