aboutsummaryrefslogtreecommitdiffhomepage
path: root/tensorflow/python/data/kernel_tests/batch_dataset_op_test.py
diff options
context:
space:
mode:
Diffstat (limited to 'tensorflow/python/data/kernel_tests/batch_dataset_op_test.py')
-rw-r--r--tensorflow/python/data/kernel_tests/batch_dataset_op_test.py4
1 files changed, 2 insertions, 2 deletions
diff --git a/tensorflow/python/data/kernel_tests/batch_dataset_op_test.py b/tensorflow/python/data/kernel_tests/batch_dataset_op_test.py
index c3d42b49af..89de55dd4f 100644
--- a/tensorflow/python/data/kernel_tests/batch_dataset_op_test.py
+++ b/tensorflow/python/data/kernel_tests/batch_dataset_op_test.py
@@ -278,7 +278,7 @@ class PaddedBatchDatasetTest(test.TestCase, parameterized.TestCase):
result = sess.run(get_next)
padded_len = padded_shapes[0]
if padded_len is None or padded_len == -1:
- padded_len = np.max(result)
+ padded_len = np.max(result) if result.size > 0 else 0
self.assertEqual((batch_size, padded_len), result.shape)
for j in range(batch_size):
seq_len = seq_lens[(i * batch_size) + j]
@@ -288,7 +288,7 @@ class PaddedBatchDatasetTest(test.TestCase, parameterized.TestCase):
if not drop_remainder and len(seq_lens) % batch_size > 0:
result = sess.run(get_next)
- padded_len = np.max(result)
+ padded_len = np.max(result) if result.size > 0 else 0
self.assertEqual((len(seq_lens) % batch_size, padded_len),
result.shape)
for j in range(len(seq_lens) % batch_size):