This is the code in the generator. It uses argmax to choose the best token at each timestep.
Now I’m starting to doubt if there is something to do with PyTorch Transformer.
def generate(self, src, src_key_padding_mask=None):
        ''' src has dimension of LEN x 1 '''
        src = self.embedding(src)
        src = self.pos_encoder(src)
        src = self.transformer_encoder(src, src_key_padding_mask=src_key_padding_mask)
        
        inputs = [self.args.sos_idx]
        for i in range(self.args.max_tgt_len):
            tgt = torch.LongTensor([inputs]).view(-1,1).to(device)
            tgt_mask = self.get_mask(i+1).to(device)
            tgt = self.embedding(tgt)
            tgt = self.pos_encoder(tgt)
            output = self.transformer_decoder(
                tgt=tgt, 
                memory=src, 
                tgt_mask=tgt_mask,
                memory_key_padding_mask = src_key_padding_mask )
            
            output = self.linear(output)
            output = self.softmax(output)
            output = output[-1] # the last timestep
            
            values, indices = output.max(dim=-1) 
            pred_token = indices.item()
            inputs.append(pred_token)

            # end when reach eos idx
            if pred_token == self.args.eos_idx:
                break