| 58 | |
| 59 | |
| 60 | class FeedForward(nn.Module): |
| 61 | def __init__(self, dim, hidden_dim, dropout_p): |
| 62 | super().__init__() |
| 63 | self.w1 = nn.Linear(dim, hidden_dim) |
| 64 | self.gelu = nn.GELU() |
| 65 | self.w2 = nn.Linear(hidden_dim, dim) |
| 66 | self.resid_dropout = nn.Dropout(dropout_p) |
| 67 | |
| 68 | def forward(self, x): |
| 69 | return self.resid_dropout(self.w2(self.gelu(self.w1(x)))) |
| 70 | |
| 71 | def reset_parameters(self): |
| 72 | self.w1.reset_parameters() |
| 73 | self.w2.reset_parameters() |
| 74 | |
| 75 | |
| 76 | class TransformerBlock(nn.Module): |