Add CompositeTensor support for DistributedIterator.

PiperOrigin-RevId: 289037261 Change-Id: Iec14cd66bcad37070c19a3a2f3bb0cc524e79ce9
bmzhao · Jan 10, 2020 · b7a6d31 · b7a6d31
1 parent 70d5a38
commit b7a6d31
Show file tree

Hide file tree

Showing 4 changed files with 125 additions and 746 deletions.
diff --git a/tensorflow/python/distribute/BUILD b/tensorflow/python/distribute/BUILD
@@ -734,10 +734,7 @@ distribute_py_test(
     name = "input_lib_test",
     srcs = ["input_lib_test.py"],
     main = "input_lib_test.py",
-    shard_count = 30,
-    tags = [
-        "no_gpu_presubmit",
-    ],
+    shard_count = 10,
     deps = [
         ":collective_all_reduce_strategy",
         ":mirrored_strategy",
@@ -1233,10 +1230,9 @@ distribute_py_test(
     name = "ctl_correctness_test",
     srcs = ["ctl_correctness_test.py"],
     main = "ctl_correctness_test.py",
-    shard_count = 30,
+    shard_count = 10,
     tags = [
         "multi_and_single_gpu",
-        "no_gpu_presubmit",
         "noguitar",  # b/140755528
     ],
     deps = [