  
      return logits, predictions
  
-@@ -325,14 +328,15 @@ def compute_metrics(logits, labels):
+@@ -325,14 +328,16 @@ def compute_metrics(logits, labels):
  
  
  @jax.jit

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

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