Fix transformer fp16 test

PiperOrigin-RevId: 263812316

Fix transformer fp16 test
PiperOrigin-RevId: 263812316
7b57b6d7 · Reed Wanderman-Milne · A. Unique TensorFlower · 07c09ccc · 7b57b6d7
Commit 7b57b6d7 authored Aug 16, 2019 by Reed Wanderman-Milne Committed by A. Unique TensorFlower Aug 16, 2019
Show whitespace changes
Inline Side-by-side

Showing with 2 additions and 2 deletions

official/transformer/v2/transformer_main_test.py official/transformer/v2/transformer_main_test.py +2 -2

No files found.
--- a/official/transformer/v2/transformer_main_test.py
+++ b/official/transformer/v2/transformer_main_test.py
@@ -103,9 +103,9 @@ class TransformerTaskTest(tf.test.TestCase):
    t = tm.TransformerTask(FLAGS)
    t.train()

-  @unittest.skipIf(tf.test.is_built_with_cuda(), 'TODO(b/139497127): '
-     'Test is skipped because tf.pad doesn\'t work with GPU.')
+  @unittest.skipUnless(tf.test.is_built_with_cuda(), 'requires GPU')
  def test_train_fp16(self):
+    FLAGS.distribution_strategy = 'one_device'
    FLAGS.dtype = 'fp16'
    t = tm.TransformerTask(FLAGS)
    t.train()