1 files changed, 4 insertions, 2 deletions
diff --git a/tensorflow/tools/dist_test/local_test.sh b/tensorflow/tools/dist_test/local_test.sh
index f9f37ff0e1..f536beef54 100755
--- a/tensorflow/tools/dist_test/local_test.sh
+++ b/tensorflow/tools/dist_test/local_test.sh
@@ -151,6 +151,8 @@ rm -rf "${BUILD_DIR}"
 # Run docker image for test.
 docker run ${DOCKER_IMG_NAME} \
     /var/tf_dist_test/scripts/dist_mnist_test.sh \
-    --ps_hosts "localhost:2000,localhost:2001" \
-    --worker_hosts "localhost:3000,localhost:3001" \
+    --ps_hosts $(seq -f "localhost:%g" -s "," \
+                 2000 $((2000 + ${NUM_PARAMETER_SERVERS} - 1))) \
+    --worker_hosts $(seq -f "localhost:%g" -s "," \
+                     3000 $((3000 + ${NUM_WORKERS} - 1))) \
     --num_gpus 0 ${SYNC_REPLICAS_FLAG}