        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):
+@@ -226,7 +228,7 @@ class Seq2seq(nn.Module):
    def apply(self,
              encoder_inputs,
              decoder_inputs,

         decode_batch(optimizer.target, 5)
   return optimizer.target
 
diff --git a/examples/seq2seq/train_test.py b/examples/seq2seq/train_test.py
index 5a8a24b..6df633e 100644
--- a/examples/seq2seq/train_test.py
+++ b/examples/seq2seq/train_test.py
@@ -102,7 +102,7 @@ class TrainTest(absltest.TestCase):
       model = train.create_model()
       optimizer = train.create_optimizer(model, 0.003)
       optimizer, train_metrics = train.train_step(
-          optimizer, batch, nn.make_rng())
+          optimizer, batch, 0.5, nn.make_rng())
 
     self.assertLessEqual(train_metrics['loss'], 5)
     self.assertGreaterEqual(train_metrics['accuracy'], 0)

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."""
 
