 +            metrics['accuracy'] * 100)
          decode_batch(optimizer.target, 5)
    return optimizer.target
+ 
+diff --git a/examples/seq2seq/train_test.py b/examples/seq2seq/train_test.py
+index b638d55..26fc7b8 100644
+--- a/examples/seq2seq/train_test.py
++++ b/examples/seq2seq/train_test.py
+@@ -101,7 +101,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)
+@@ -113,3 +113,4 @@ class TrainTest(absltest.TestCase):
+ 
+ if __name__ == '__main__':
+   absltest.main()
++

       hidden_size: int, the number of hidden dimensions in the encoder and
         decoder LSTMs.
     Returns:
@@ -252,8 +255,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)
+    logits, predictions = Decoder(init_decoder_state, decoder_inputs[:, :-1])
 
     return logits, predictions
 
@@ -323,14 +325,16 @@ def compute_metrics(logits, labels):
 
 
 @jax.jit

     loss = cross_entropy_loss(logits, labels, get_sequence_lengths(labels))
     return loss, logits
   grad_fn = jax.value_and_grad(loss_fn, has_aux=True)
@@ -354,7 +359,7 @@ def decode(model, inputs, rng):
   init_decoder_inputs = jnp.tile(init_decoder_input,
                                  (inputs.shape[0], get_max_output_len(), 1))
   with nn.stochastic(rng):
