| 598 | |
| 599 | |
| 600 | class Block(BaseBlock): |
| 601 | def __init__(self, config: Config) -> None: |
| 602 | nn.Module.__init__(self) |
| 603 | self.norm_1 = config.norm_class(config.n_embd, eps=config.norm_eps) |
| 604 | self.attn = CausalSelfAttention(config) |
| 605 | if not config.shared_attention_norm: |
| 606 | self.norm_2 = config.norm_class(config.n_embd, eps=config.norm_eps) |
| 607 | self.mlp = config.mlp_class(config) |
| 608 | |
| 609 | self.config = config |
| 610 | |
| 611 | |
| 612 | class CausalSelfAttention(BaseCausalSelfAttention): |