Method __init__

src/shape_assembly/models/train/transformer.py:20–26

(self, encoder, decoder, src_embed, tgt_embed, generator)

class EncoderDecoder(pl.LightningModule):
    def __init__(self, encoder, decoder, src_embed, tgt_embed, generator):
        super().__init__()
        self.encoder = encoder
        self.decoder = decoder
        self.src_embed = src_embed
        self.tgt_embed = tgt_embed
        self.generator = generator

    def encode(self, src, src_mask):
        return self.encoder(self.src_embed(src), src_mask)