The implementation is identical to `lit_gpt.model.Block` with the exception that we replace the attention layer where adaption is implemented.
| 111 | |
| 112 | |
| 113 | class Block(BaseBlock): |
| 114 | """The implementation is identical to `lit_gpt.model.Block` with the exception that |
| 115 | we replace the attention layer where adaption is implemented.""" |
| 116 | |
| 117 | def __init__(self, config: Config, block_idx: int) -> None: |
| 118 | # Skip the parent class __init__ altogether and replace it to avoid useless allocations |
| 119 | nn.Module.__init__(self) |
| 120 | self.norm_1 = config.norm_class(config.n_embd, eps=config.norm_eps) |
| 121 | self.attn = CausalSelfAttention(config, block_idx) |
| 122 | if not config.shared_attention_norm: |
| 123 | self.norm_2 = config.norm_class(config.n_embd, eps=config.norm_eps) |
| 124 | self.mlp = config.mlp_class(config) |
| 125 | |
| 126 | self.config = config |
| 127 | |
| 128 | |
| 129 | class CausalSelfAttention(BaseCausalSelfAttention): |