I have implemented simple model based on nn.Transformer and trained on Multi30k dataset. The final model always predicts the same output for different inputs. Can you please suggest what is wrong?
https://colab.research.google.com/drive/1KssnWND4dKx_aBTW0TQZlI4THttHiFTm?usp=sharing 2
class TransformerBase(nn.Module):
    def __init__(self, input_size, output_size, hidden_size=256):
        super(TransformerBase, self).__init__()
        
        self.input_size = input_size
        self.output_size = output_size
        self.hidden_size = hidden_size
              
        self.transformer_model = nn.Transformer(d_model=hidden_size)
        self.embedding_input = nn.Embedding(self.input_size, hidden_size)
        self.embedding_output = nn.Embedding(self.output_size, hidden_size)
        self.pos_encoder = PositionalEncoding(hidden_size, 0.1)
        self.fc_out = nn.Linear(hidden_size, self.output_size)
        
    def forward(self, src, trg):
        embedded_input = self.pos_encoder(self.embedding_input(src) * math.sqrt(self.hidden_size))
        embedding_output = self.pos_encoder(self.embedding_output(trg) * math.sqrt(self.hidden_size))
        
        tgt_mask = self.transformer_model.generate_square_subsequent_mask(trg.shape[0])
        x = self.transformer_model(src=embedded_input, tgt=embedding_output, tgt_mask=tgt_mask.to(device))
        
        return self.fc_out(x)