Fix transformer test (#4606)

b1a704d7 · Katherine Wu · GitHub · c39a88c0 · b1a704d7
Unverified Commit b1a704d7 authored Jun 22, 2018 by Katherine Wu Committed by GitHub Jun 22, 2018
Show whitespace changes
Inline Side-by-side

Showing with 2 additions and 2 deletions

official/transformer/model/beam_search_test.py official/transformer/model/beam_search_test.py +2 -2

No files found.
--- a/official/transformer/model/beam_search_test.py
+++ b/official/transformer/model/beam_search_test.py
@@ -33,8 +33,8 @@ class BeamSearchHelperTests(tf.test.TestCase):
    self.assertAllEqual([7, 3, 4, 2, 5], shape)

  def test_shape_list(self):
-    y = tf.constant(4.0)
-    x = tf.ones([7, tf.to_int32(tf.sqrt(y)), 2, 5])
+    y = tf.placeholder(dtype=tf.int32, shape=[])
+    x = tf.ones([7, y, 2, 5])
    shape = beam_search._shape_list(x)
    self.assertIsInstance(shape[0], int)
    self.assertIsInstance(shape[1], tf.Tensor)