| 89 | return x |
| 90 | |
| 91 | class Decoder(nn.Module): |
| 92 | def __init__(self, embed_dim, n_blocks, n_heads, dropout): |
| 93 | super(Decoder, self).__init__() |
| 94 | self.layers = nn.ModuleList([DecoderBlock(embed_dim, n_heads, dropout) for _ in range(n_blocks)]) |
| 95 | |
| 96 | def forward(self, x, enc_output, tgt_mask=None, memory_mask=None): |
| 97 | for layer in self.layers: |
| 98 | x = layer(x, enc_output, tgt_mask, memory_mask) |
| 99 | return x |
| 100 | |
| 101 | class Transformer(nn.Module): |
| 102 | def __init__(self, input_len, embed_dim, n_blocks, n_heads, dropout): |