diff --git a/mesh_tensorflow/transformer/transformer.py b/mesh_tensorflow/transformer/transformer.py
index <HASH>..<HASH> 100644
--- a/mesh_tensorflow/transformer/transformer.py
+++ b/mesh_tensorflow/transformer/transformer.py
@@ -997,7 +997,7 @@ class Unitransformer(object):
         mesh_shape=self.mesh_shape,
         layout=self.layout)
     return mtf.gather(
-        beams, mtf.constant(inputs.mesh, 2, dtype=tf.int32), beam_dim)
+        beams, mtf.constant(inputs.mesh, 0, dtype=tf.int32), beam_dim)
 
 
 @gin.configurable
Revert a mistake I made in cl/<I> . This should fix beam search.

PiperOrigin-RevId: <I>