[tf.data] Scaling down the `batch_dataset_op_test`.

PiperOrigin-RevId: 198715407

diff --git a/tensorflow/contrib/data/python/kernel_tests/BUILD b/tensorflow/contrib/data/python/kernel_tests/BUILD
index 285c77d..c483a43 100644 (file)
--- a/tensorflow/contrib/data/python/kernel_tests/BUILD
+++ b/tensorflow/contrib/data/python/kernel_tests/BUILD
@@ -8,7 +8,7 @@ load("//tensorflow:tensorflow.bzl", "cuda_py_test", "py_test", "tf_py_test")
 
 py_test(
     name = "batch_dataset_op_test",
-    size = "large",
+    size = "medium",
     srcs = ["batch_dataset_op_test.py"],
     srcs_version = "PY2AND3",
     tags = [

diff --git a/tensorflow/contrib/data/python/kernel_tests/batch_dataset_op_test.py b/tensorflow/contrib/data/python/kernel_tests/batch_dataset_op_test.py
index e309d61..b5fbc45 100644 (file)
--- a/tensorflow/contrib/data/python/kernel_tests/batch_dataset_op_test.py
+++ b/tensorflow/contrib/data/python/kernel_tests/batch_dataset_op_test.py
@@ -553,14 +553,14 @@ class BatchDatasetTest(test.TestCase):
         sess.run(next_element)
 
   def testMapAndBatchParallelGetNext(self):
-    iterator = (dataset_ops.Dataset.range(500000)
+    iterator = (dataset_ops.Dataset.range(50000)
                 .apply(batching.map_and_batch(lambda x: x, batch_size=100))
                 .make_one_shot_iterator())
     elements = []
     for _ in range(100):
       elements.append(iterator.get_next())
     with self.test_session() as sess:
-      for i in range(50):
+      for i in range(5):
         got = sess.run(elements)
         got.sort(key=lambda x: x[0])
         expected = []
@@ -572,7 +572,7 @@ class BatchDatasetTest(test.TestCase):
 
   def testMapAndBatchParallelGetNextDropRemainder(self):
     iterator = (
-        dataset_ops.Dataset.range(499999).apply(
+        dataset_ops.Dataset.range(49999).apply(
             batching.map_and_batch(
                 lambda x: x, batch_size=100, drop_remainder=True))
         .make_one_shot_iterator())
@@ -580,7 +580,7 @@ class BatchDatasetTest(test.TestCase):
     for _ in range(100):
       elements.append(iterator.get_next())
     with self.test_session() as sess:
-      for i in range(49):
+      for i in range(4):
         got = sess.run(elements)
         got.sort(key=lambda x: x[0])
         expected = []