def __init__()

in summarize_from_feedback/models/transformer.py [0:0]


    def __init__(
        self,
        n_ctx,
        n_vocab,
        d_model=128,
        n_layer=2,
        heads=1,
        attn_dropout=0.0,
        resid_dropout=0.0,
        emb_dropout=0.0,
        zero_out=False,
        init_scale=1.0,
        res_scale=False,
        m_attn=0.25,
        m_mlp=1,
        mp_comm: Comm = None,
        include_pos_embeddings=True,
        include_input_embeddings=True,
        include_output_unembeddings=True,
        # For e.g. reward model training, we want the final layer norm before the extra head
        # but not the output unembeddings, so control this separately.
        include_final_layer_norm=True,
        afn="quick_gelu",
        key_bias=False,
        flatten_multi_index_batch_dims=False,  # allows the first dims to be batch dims and flattens them, outputs will be flattened
        global_idxs_for_resblocks: Optional[List[int]] = None,
        **extra_args,