Method forward

src/shape_assembly/models/train/transformer.py:34–35

(self, src, tgt, src_mask, tgt_mask)

        return self.generator(self.decoder(self.tgt_embed(tgt), memory, src_mask, tgt_mask))

    def forward(self, src, tgt, src_mask, tgt_mask):
        return self.decode(self.encode(src, src_mask), src_mask, tgt, tgt_mask)

class Encoder(pl.LightningModule):
    def __init__(self, layer, N):