Fix multi_worker_callback_tf2_test test target by only running it wit…

…h CPU. The test is not using GPU anyway. PiperOrigin-RevId: 294784694 Change-Id: I9e9d2f8db05160799ef64b43f0cc8b1a927637e0
gaurav1086 · Feb 13, 2020 · fc36231 · fc36231
1 parent 4f2afc0
commit fc36231
Showing 1 changed file with 2 additions and 5 deletions.
diff --git a/tensorflow/python/keras/distribute/BUILD b/tensorflow/python/keras/distribute/BUILD
@@ -374,14 +374,11 @@ cuda_py_test(
     ],
 )
 
-cuda_py_test(
+py_test(
     name = "multi_worker_callback_tf2_test",
     srcs = ["multi_worker_callback_tf2_test.py"],
+    python_version = "PY3",
     shard_count = 10,
-    tags = [
-        "manual",  # TODO(b/148795660): Reenable once the flakiness is fixed.
-        "notap",
-    ],
     deps = [
         "//tensorflow/python/distribute:collective_all_reduce_strategy",
         "//tensorflow/python/distribute:combinations",