        hidden_size: int, the number of hidden dimensions in the encoder and
          decoder LSTMs.
      Returns:
-@@ -265,7 +268,7 @@ class Seq2seq(nn.Module):
-     logits, predictions = decoder(
-         init_decoder_state,
-         decoder_inputs[:, :-1],
+@@ -253,7 +256,7 @@ class Seq2seq(nn.Module):
+     init_decoder_state = Encoder(encoder_inputs, eos_id=eos_id,
+         hidden_size=hidden_size)
+     logits, predictions = Decoder(init_decoder_state, decoder_inputs[:, :-1], 
 -        teacher_force=teacher_force)
-+        sample_probability=sample_probability)
++        teacher_force=teacher_force, sample_probability=sample_probability)
  
      return logits, predictions
  
-@@ -335,14 +338,16 @@ def compute_metrics(logits, labels):
+@@ -323,14 +326,16 @@ def compute_metrics(logits, labels):
  
  
  @jax.jit

       lstm_state, y = lstm_cell(lstm_state, x)
       logits = projection(y)
       predicted_tokens = jax.random.categorical(categorical_rng, logits)
@@ -233,7 +235,7 @@ class Seq2seq(nn.Module):
   def apply(self,
             encoder_inputs,
             decoder_inputs,

diff --git a/examples/seq2seq/train.py b/examples/seq2seq/train.py
index 60b79b5..c3df5b1 100644
--- a/examples/seq2seq/train.py
+++ b/examples/seq2seq/train.py
@@ -190,17 +190,19 @@ class Encoder(nn.Module):
 class Decoder(nn.Module):
   """LSTM decoder."""
 
