 diff --git a/examples/seq2seq/train.py b/examples/seq2seq/train.py
-index 9b64f88..3eba1db 100644
+index 9b64f88..69bed4f 100644
 --- a/examples/seq2seq/train.py
 +++ b/examples/seq2seq/train.py
 @@ -186,17 +186,19 @@ class Encoder(nn.Module):

         decode_batch(optimizer.target, 5)
   return optimizer.target
 

   return predictions
 
 
@@ -374,15 +379,22 @@ def decode_batch(model, batch_size):
 
 def train_model():
   """Train for a fixed number of steps and decode during training."""
