| 77 | return x + sublayer(self.norm(x)) |
| 78 | |
| 79 | class EncoderLayer(pl.LightningModule): |
| 80 | def __init__(self, size, self_attn, feed_forward): |
| 81 | super(EncoderLayer, self).__init__() |
| 82 | self.self_attn = self_attn |
| 83 | self.feed_forward = feed_forward |
| 84 | self.sublayer = clones(SublayerConnection(size), 2) |
| 85 | self.size = size |
| 86 | |
| 87 | def forward(self, x, mask): |
| 88 | x = self.sublayer[0](x, lambda x: self.self_attn(x, x, x, mask)) |
| 89 | return self.sublayer[1](x, self.feed_forward) |
| 90 | |
| 91 | class DecoderLayer(pl.LightningModule): |
| 92 |