| 14 | |
| 15 | |
| 16 | class Encoder(nn.Module): |
| 17 | def __init__(self, d_model, N, heads, dropout): |
| 18 | super().__init__() |
| 19 | self.N = N |
| 20 | self.layers = get_clones(Layer(d_model, heads, dropout), N) |
| 21 | self.norm = Norm(d_model) |
| 22 | |
| 23 | def forward(self, src, mask=None): |
| 24 | x = src |
| 25 | for i in range(self.N): |
| 26 | x = self.layers[i](x, mask) |
| 27 | return self.norm(x) |
| 28 | |
| 29 | |
| 30 | class AE(nn.Module): |