Method __init__

src/shape_assembly/models/train/transformer.py:80–85

(self, size, self_attn, feed_forward)

class EncoderLayer(pl.LightningModule):
    def __init__(self, size, self_attn, feed_forward):
        super(EncoderLayer, self).__init__()
        self.self_attn = self_attn
        self.feed_forward = feed_forward
        self.sublayer = clones(SublayerConnection(size), 2)
        self.size = size

    def forward(self, x, mask):
        x = self.sublayer[0](x, lambda x: self.self_attn(x, x, x, mask))