From b100b185eecacef9990525e9a712b5547fa20689 Mon Sep 17 00:00:00 2001 From: "A. Unique TensorFlower" Date: Tue, 26 May 2020 17:28:20 -0700 Subject: [PATCH] Don't crash in 3D pooling ops with empty batch size on GPU. PiperOrigin-RevId: 313299099 Change-Id: I40ce8f57efc386ae820460a325cfebee1be14d77 --- tensorflow/core/kernels/pooling_ops_3d.cc | 1 + tensorflow/python/kernel_tests/pooling_ops_3d_test.py | 4 ++-- 2 files changed, 3 insertions(+), 2 deletions(-) diff --git a/tensorflow/core/kernels/pooling_ops_3d.cc b/tensorflow/core/kernels/pooling_ops_3d.cc index 31ead11dd34..532d861e615 100644 --- a/tensorflow/core/kernels/pooling_ops_3d.cc +++ b/tensorflow/core/kernels/pooling_ops_3d.cc @@ -192,6 +192,7 @@ class Pooling3DOp : public UnaryOp { {{out[2], out[1], out[0]}}, depth); Tensor* output; OP_REQUIRES_OK(context, context->allocate_output(0, out_shape, &output)); + if (out_shape.num_elements() == 0) return; LaunchPoolingOp::launch(context, tensor_in, window, stride, padding, data_format_, padding_, output); diff --git a/tensorflow/python/kernel_tests/pooling_ops_3d_test.py b/tensorflow/python/kernel_tests/pooling_ops_3d_test.py index d5331dcb3e9..051f7e1168a 100644 --- a/tensorflow/python/kernel_tests/pooling_ops_3d_test.py +++ b/tensorflow/python/kernel_tests/pooling_ops_3d_test.py @@ -205,14 +205,14 @@ class PoolingTest(test.TestCase): padding="VALID", expected=[29.5, 32.5, 50.5, 53.5, 176.5, 179.5, 197.5, 200.5]) - def _MaxPool3DEmptyTensorOutputShape(self): + def testMaxPool3DEmptyTensorOutputShape(self): """Verifies the output shape of the max pooling function when tensor is empty. Args: none """ input_sizes = [0, 112, 112, 112, 64] - input_data = 1 + input_data = 1. input_tensor = constant_op.constant( input_data, shape=input_sizes, name="input") max_pool_3d = nn_ops.max_pool3d(