Class Config

axlearn/common/bert.py:215–217

    @config_class
    class Config(BaseClassificationHead.Config):
        pooler: BertPooler.Config = BertPooler.default_config()  # Pooler layer.
        output: Linear.Config = Linear.default_config()  # Output projection layer.

    def __init__(self, cfg: Config, *, parent: Optional[Module]):
        super().__init__(cfg, parent=parent)