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