(self, args: ModelArgs)
| 245 | |
| 246 | class Transformer(nn.Module): |
| 247 | def __init__(self, args: ModelArgs): |
| 248 | super().__init__() |
| 249 | assert args.vocab_size > 0 |
| 250 | |
| 251 | self.tok_embeddings = nn.Embedding( |
| 252 | num_embeddings=args.vocab_size, |
| 253 | embedding_dim=args.dim, |
| 254 | ) |
| 255 | |
| 256 | self.layers = nn.ModuleList() |
| 257 | for _ in range(args.n_layers): |
| 258 | self.layers.append(TransformerBlock(args)) |
| 259 | |
| 260 | self.norm = RMSNorm(args.dim, eps=args.norm_eps) |
| 261 | |
| 262 | self.output = nn.Linear( |
| 263 | args.dim, |
| 264 | args.vocab_size, |
| 265 | bias=False, |
| 266 | ) |
| 267 | |
| 268 | @torch.no_grad() |
| 269 | def forward_with_attn_bias( |
nothing calls this directly
no test coverage detected