| 393 | return x |
| 394 | |
| 395 | class EncoderLayer(nn.Module): |
| 396 | def __init__(self, d_input, d_model, heads, dropout=0.1): |
| 397 | super().__init__() |
| 398 | self.input_linear = nn.Linear(d_input, d_model) |
| 399 | self.norm_1 = Norm(d_model) |
| 400 | self.norm_2 = Norm(d_model) |
| 401 | self.attn = MultiHeadAttention(heads, d_model, dropout=dropout) |
| 402 | self.ff = FeedForward(d_model, dropout=dropout) |
| 403 | self.dropout_1 = nn.Dropout(dropout) |
| 404 | self.dropout_2 = nn.Dropout(dropout) |
| 405 | |
| 406 | def forward(self, x, mask=None): |
| 407 | x2 = self.norm_1(x) |
| 408 | x = x + self.dropout_1(self.attn(x2, x2, x2, mask)) |
| 409 | x2 = self.norm_2(x) |
| 410 | x = x + self.dropout_2(self.ff(x2)) |
| 411 | return x |
| 412 | |
| 413 | |
| 414 | # build a decoder layer with two multi-head attention layers and |