        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,

diff --git a/examples/seq2seq/train.py b/examples/seq2seq/train.py
index 1597079..77449b7 100644
--- a/examples/seq2seq/train.py
+++ b/examples/seq2seq/train.py
@@ -191,17 +191,19 @@ class Encoder(nn.Module):
 class Decoder(nn.Module):
   """LSTM decoder."""
 

     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):
