| 562 | # Mapping network |
| 563 | |
| 564 | class MappingFeedForwardBlock(nn.Module): |
| 565 | def __init__(self, d_model, d_ff, dropout=0.0): |
| 566 | super().__init__() |
| 567 | self.norm = RMSNorm(d_model) |
| 568 | self.up_proj = apply_wd(LinearGEGLU(d_model, d_ff, bias=False)) |
| 569 | self.dropout = nn.Dropout(dropout) |
| 570 | self.down_proj = apply_wd(zero_init(Linear(d_ff, d_model, bias=False))) |
| 571 | |
| 572 | def forward(self, x): |
| 573 | skip = x |
| 574 | x = self.norm(x) |
| 575 | x = self.up_proj(x) |
| 576 | x = self.dropout(x) |
| 577 | x = self.down_proj(x) |
| 578 | return x + skip |
| 579 | |
| 580 | class MappingNetwork(nn.Module): |
| 581 | def __init__(self, n_layers, d_model, d_ff, dropout=0.0): |