Yeah @mathematicsofpaul , I doubt my model will work for a real dataset without it. This is how I implemented my forward function. I’m going to spend tomorrow looking more closely at the PositionalEncoding
def forward(self, src, tgt):

    # create padding masks
    src_key_padding_mask = (src == self.src_vocab.stoi['<pad>'])  # N S
    tgt_key_padding_mask = (tgt == self.tgt_vocab.stoi['<pad>'])  # N T
    memory_key_padding_mask = src_key_padding_mask.clone()        # N S

    # prevent attention of future output tokens
    tgt_mask = self.transformer.generate_square_subsequent_mask(tgt.shape[-1])  # T T

    src = rearrange(src, 'n s -> s n')
    tgt = rearrange(tgt, 'n t -> t n')

    src = self.embed_src(src) * math.sqrt(self.d_model) # S, N, E
    tgt = self.embed_tgt(tgt) * math.sqrt(self.d_model) # T, N, E

    output = self.transformer(
        src, tgt, 
        tgt_mask=tgt_mask, 
        src_key_padding_mask=src_key_padding_mask,
        tgt_key_padding_mask=tgt_key_padding_mask, 

    output = rearrange(output, 't n e -> n t e')
    return self.fc(output)

and my train and validation are the same
        # reshape inputs
        src = rearrange(batch.object_name, 's n -> n s')
        tgt = rearrange(batch.tags, 't n -> n t')   

        # Create tgt_inp and tgt_out (which is tgt_inp but shifted by 1)
        tgt_inp, tgt_out = tgt[:,:-1], tgt[:,1:]

        optim.zero_grad()
        out = model(src, tgt_inp)
        #predicted = F.log_softmax(out, dim=-1).argmax(dim=-1)
        loss = criterion(rearrange(out, 'b t v -> (b t) v'), rearrange(tgt_out, 'b o -> (b o)'))

        # Backpropagate and update optim
        loss.backward()
        optim.step()