| 89 | return self.sublayer[1](x, self.feed_forward) |
| 90 | |
| 91 | class DecoderLayer(pl.LightningModule): |
| 92 | |
| 93 | def __init__(self, size, self_attn, src_attn, feed_forward): |
| 94 | super().__init__() |
| 95 | self.size = size |
| 96 | self.self_attn = self_attn |
| 97 | self.src_attn = src_attn |
| 98 | self.feed_forward = feed_forward |
| 99 | self.sublayer = clones(SublayerConnection(size), 3) |
| 100 | |
| 101 | def forward(self, x, memory, src_mask, tgt_mask): |
| 102 | m = memory |
| 103 | x = self.sublayer[0](x, lambda x: self.self_attn(x, x, x, tgt_mask)) |
| 104 | x = self.sublayer[1](x, lambda x: self.src_attn(x, m, m, src_mask)) |
| 105 | return self.sublayer[2](x, self.feed_forward) |
| 106 | |
| 107 | class MultiHeadedAttention(pl.LightningModule): |
| 108 | |