typo

Author: tcapelle

model.py

@@ -223,7 +223,7 @@ def project(self, x):
         # (batch, seq_len, vocab_size)
         return self.projection_layer(x)
 
-    def forward(self, encoder_output, encoder_mask, decoder_input, decoder_mask):      
+    def forward(self, encoder_input, encoder_mask, decoder_input, decoder_mask):      
         encoder_output = self.encode(encoder_input, encoder_mask) # (B, seq_len, d_model)
         decoder_output = self.decode(encoder_output, encoder_mask, decoder_input, decoder_mask) # (B, seq_len, d_model)
         return self.project(decoder_output) # (B, seq_len, vocab_size))

train.py

@@ -265,8 +265,8 @@ def train_model(config: ModelConfig):
             # encoder_output = model.module.encode(encoder_input, encoder_mask) # (B, seq_len, d_model)
             # decoder_output = model.module.decode(encoder_output, encoder_mask, decoder_input, decoder_mask) # (B, seq_len, d_model)
             # proj_output = model.module.project(decoder_output) # (B, seq_len, vocab_size)
-            proj_output = model(encoder_output, encoder_mask, decoder_input, decoder_mask)
-            
+            proj_output = model(encoder_input, encoder_mask, decoder_input, decoder_mask)
+
             # Compare the output with the label
             label = batch['label'].to(device) # (B, seq_len)