| 35 | return self.decode(self.encode(src, src_mask), src_mask, tgt, tgt_mask) |
| 36 | |
| 37 | class Encoder(pl.LightningModule): |
| 38 | def __init__(self, layer, N): |
| 39 | super().__init__() |
| 40 | self.layers = clones(layer, N) |
| 41 | self.norm = LayerNorm(layer.size) |
| 42 | |
| 43 | def forward(self, x, mask): |
| 44 | for layer in self.layers: |
| 45 | x = layer(x, mask) |
| 46 | return self.norm(x) |
| 47 | |
| 48 | class Decoder(pl.LightningModule): |
| 49 | def __init__(self, layer, N): |