diff --git a/tensorflow/core/kernels/in_topk_op_gpu.cu.cc b/tensorflow/core/kernels/in_topk_op_gpu.cu.cc index 44c96f67b26..f701071cb8e 100644 --- a/tensorflow/core/kernels/in_topk_op_gpu.cu.cc +++ b/tensorflow/core/kernels/in_topk_op_gpu.cu.cc @@ -100,6 +100,12 @@ struct InTopKFunctor { errors::InvalidArgument( "Number of targets * number of classes must be less than INT_MAX")); + if (num_targets == 0 || num_classes == 0) { + // Result is empty, so shortcut the rest of the function to avoid + // launching kernels with empty input. + return; + } + // Temporary storage for a mask computed by `ComputePredictionMaskKernel`. Tensor predictions_mask; OP_REQUIRES_OK( diff --git a/tensorflow/python/kernel_tests/in_topk_op_test.py b/tensorflow/python/kernel_tests/in_topk_op_test.py index 2ef233f8d68..c636cee0dd5 100644 --- a/tensorflow/python/kernel_tests/in_topk_op_test.py +++ b/tensorflow/python/kernel_tests/in_topk_op_test.py @@ -28,7 +28,7 @@ from tensorflow.python.platform import test class InTopKTest(test.TestCase): def _validateInTopK(self, predictions, target, k, expected): - np_ans = np.array(expected) + np_ans = np.array(expected, np.bool) with self.cached_session(use_gpu=True): precision = nn_ops.in_top_k(predictions, target, k) out = self.evaluate(precision) @@ -66,6 +66,11 @@ class InTopKTest(test.TestCase): target = [2, 4] # must return False for invalid target self._validateInTopK(predictions, target, 2, [True, False]) + def testEmpty(self): + predictions = np.empty([0, 5]) + target = np.empty([0], np.int32) + self._validateInTopK(predictions, target, 2, []) + def testTensorK(self): predictions = [[0.1, 0.3, 0.2, 0.4], [0.1, 0.2, 0.3, 0.4]] target = [0, 2]