(self, x: torch.Tensor)
| 190 | self.ffn_sub_norm = RMSNorm(hidden_dim, norm_eps) |
| 191 | |
| 192 | def forward(self, x: torch.Tensor) -> torch.Tensor: |
| 193 | x13 = self.w13(x) |
| 194 | x1, x3 = x13.chunk(2, -1) |
| 195 | inner = self.ffn_sub_norm(squared_relu(x1) * x3) |
| 196 | output = self.w2(inner) |
| 197 | return output |
| 198 | |
| 199 | |
| 200 | class TransformerBlock(nn.Module): |
nothing calls this directly
no test coverage detected