(self, args: ModelArgs)
| 75 | |
| 76 | class TransformerBlock(nn.Module): |
| 77 | def __init__(self, args: ModelArgs): |
| 78 | super().__init__() |
| 79 | self.attention_norm = nn.LayerNorm(args.dim) |
| 80 | self.attention = Attention(args) |
| 81 | self.ffn_norm = nn.LayerNorm(args.dim) |
| 82 | self.feed_forward = FeedForward( |
| 83 | args.dim, hidden_dim=4 * args.dim, dropout_p=args.dropout_p |
| 84 | ) |
| 85 | |
| 86 | def forward(self, x): |
| 87 | h = x + self.attention(self.attention_norm(x)) |
nothing calls this directly
no test coverage detected