(
self,
dim: int,
depth: int,
num_tokens: int,
heads: int = 8,
ff_mult: int = 4,
)
| 116 | """ |
| 117 | |
| 118 | def __init__( |
| 119 | self, |
| 120 | dim: int, |
| 121 | depth: int, |
| 122 | num_tokens: int, |
| 123 | heads: int = 8, |
| 124 | ff_mult: int = 4, |
| 125 | ): |
| 126 | super().__init__() |
| 127 | self.emb = nn.Embedding(num_tokens, dim) |
| 128 | |
| 129 | self.transformer = Transformer( |
| 130 | dim=dim, depth=depth, heads=heads, ff_mult=ff_mult |
| 131 | ) |
| 132 | |
| 133 | # self.to_logits = nn.Sequential(RMSNorm(dim), nn.Linear(dim, num_tokens)) |
| 134 | self.to_logits = OutputHead( |
| 135 | dim, |
| 136 | vocab_size=num_tokens, |
| 137 | ) |
| 138 | |
| 139 | # Norm |
| 140 | self.norm = nn.LayerNorm(dim) |
| 141 | |
| 142 | def forward(self, x): |
| 143 | x = self.emb(x) |
nothing calls this directly
no test coverage detected