| 71 | |
| 72 | |
| 73 | class Block(nn.Module): |
| 74 | def __init__(self, n_embd: int, n_head: int, block_size: int, dropout: float = 0.0): |
| 75 | super().__init__() |
| 76 | self.ln_1 = nn.LayerNorm(n_embd) |
| 77 | self.attn = CausalSelfAttention(n_embd, n_head, block_size, dropout) |
| 78 | self.ln_2 = nn.LayerNorm(n_embd) |
| 79 | self.mlp = MLP(n_embd, dropout) |
| 80 | |
| 81 | def forward(self, x: torch.Tensor) -> torch.Tensor: |
| 82 | x = x + self.attn(self.ln_1(x)) |
| 83 | x = x + self.mlp(self.ln_2(x)) |
| 84 | return x |
| 85 | |
| 86 | |
| 87 | class GPT2(nn.Module): |