commit 4e08364f85b79ae9bdeb7785a94ed5c857672876
author Rick Chao <rchao@google.com> Thu Apr 01 14:40:18 2021 -0700
committer TensorFlower Gardener <gardener@tensorflow.org> Thu Apr 01 14:43:50 2021 -0700
tree 918e8a700b71ee7b7db7ca8e1246d5d7223cbcaf
parent 68c4ed4e40133d04526a43d9820be9ec89fa1131

PSv2: Increase the shard count for parameter server training test to avoid a timeout and unblock tap.

PiperOrigin-RevId: 366333433
Change-Id: I790140625a4c2f7dcf35920e642f4de75e36cad2

tensorflow/python/keras/distribute/BUILD

diff --git a/tensorflow/python/keras/distribute/BUILD b/tensorflow/python/keras/distribute/BUILD
index 55ed562..0a0453b 100644
--- a/tensorflow/python/keras/distribute/BUILD
+++ b/tensorflow/python/keras/distribute/BUILD
@@ -860,7 +860,7 @@
     name = "parameter_server_training_test",
     srcs = ["parameter_server_training_test.py"],
     python_version = "PY3",
-    shard_count = 1,
+    shard_count = 4,  # TODO(b/184290570): Investigate why only 1 shard times out.
     tags = [
         "multi_and_single_gpu",
         "no_tfrt",  # TODO(b/180537361): Reenable TFRT after the issue is resolved.