Internal change

PiperOrigin-RevId: 283622345

Internal change
PiperOrigin-RevId: 283622345
1d3cd3cf · Hongkun Yu · A. Unique TensorFlower · b478430d · 1d3cd3cf
Commit 1d3cd3cf authored Dec 03, 2019 by Hongkun Yu Committed by A. Unique TensorFlower Dec 03, 2019
Hide whitespace changes
Inline Side-by-side

Showing with 3 additions and 1 deletion

official/transformer/model/beam_search.py official/transformer/model/beam_search.py +3 -1

No files found.
--- a/official/transformer/model/beam_search.py
+++ b/official/transformer/model/beam_search.py
@@ -402,7 +402,9 @@ class SequenceBeamSearch(object):
    topk_ids = topk_indices % self.vocab_size
    if self.padded_decode:
      topk_seq = tf.transpose(topk_seq, perm=[2, 0, 1])
-      topk_seq = tf.tensor_scatter_nd_update(topk_seq, [i + 1], topk_ids)
+      # TODO(b/145533236, hongkuny): Reverts once TF fix the validation.
+      topk_seq = tf.tensor_scatter_nd_update(topk_seq, [[i + 1]],
+                                             tf.expand_dims(topk_ids, axis=0))
      topk_seq = tf.transpose(topk_seq, perm=[1, 2, 0])
    else:
      topk_ids = tf.expand_dims(topk_ids, axis=2)