| 55 | return x |
| 56 | |
| 57 | class Encoder(nn.Module): |
| 58 | def __init__(self, embed_dim, n_blocks, n_heads, dropout): |
| 59 | super(Encoder, self).__init__() |
| 60 | self.layers = nn.ModuleList([EncoderBlock(embed_dim, n_heads, dropout) for _ in range(n_blocks)]) |
| 61 | |
| 62 | def forward(self, x, src_mask=None): |
| 63 | for layer in self.layers: |
| 64 | x = layer(x, src_mask) |
| 65 | return x |
| 66 | |
| 67 | class DecoderBlock(nn.Module): |
| 68 | def __init__(self, embed_dim, n_heads, dropout): |