| 46 | return self.norm(x) |
| 47 | |
| 48 | class Decoder(pl.LightningModule): |
| 49 | def __init__(self, layer, N): |
| 50 | super().__init__() |
| 51 | self.layers = clones(layer, N) |
| 52 | self.norm = LayerNorm(layer.size) |
| 53 | |
| 54 | def forward(self, x, memory, src_mask, tgt_mask): |
| 55 | for layer in self.layers: |
| 56 | x = layer(x, memory, src_mask, tgt_mask) |
| 57 | return self.norm(x) |
| 58 | |
| 59 | class LayerNorm(pl.LightningModule): |
| 60 | def __init__(self, features, eps=1e-6): |