From 9b5f9aabd84b0d1db782286143c560792d711449 Mon Sep 17 00:00:00 2001 From: Reed Wanderman-Milne Date: Thu, 3 Oct 2019 21:33:53 -0700 Subject: [PATCH] Run bincount_op_test on GPUs. Resolves #32755 PiperOrigin-RevId: 272804557 --- tensorflow/python/kernel_tests/BUILD | 2 +- tensorflow/python/kernel_tests/bincount_op_test.py | 3 ++- 2 files changed, 3 insertions(+), 2 deletions(-) diff --git a/tensorflow/python/kernel_tests/BUILD b/tensorflow/python/kernel_tests/BUILD index aefe1a7e03b..eba090d409f 100644 --- a/tensorflow/python/kernel_tests/BUILD +++ b/tensorflow/python/kernel_tests/BUILD @@ -167,7 +167,7 @@ cuda_py_test( ], ) -tf_py_test( +cuda_py_test( name = "bincount_op_test", size = "small", srcs = ["bincount_op_test.py"], diff --git a/tensorflow/python/kernel_tests/bincount_op_test.py b/tensorflow/python/kernel_tests/bincount_op_test.py index 2168206c3a5..c77069c3cb0 100644 --- a/tensorflow/python/kernel_tests/bincount_op_test.py +++ b/tensorflow/python/kernel_tests/bincount_op_test.py @@ -21,6 +21,7 @@ import numpy as np from tensorflow.python.framework import dtypes from tensorflow.python.framework import errors +from tensorflow.python.framework import ops from tensorflow.python.framework import test_util from tensorflow.python.ops import array_ops from tensorflow.python.ops import gen_math_ops @@ -105,7 +106,7 @@ class BincountTest(test_util.TensorFlowTestCase): def test_negative(self): # unsorted_segment_sum will only report InvalidArgumentError on CPU - with self.cached_session(): + with self.cached_session(), ops.device("/CPU:0"): with self.assertRaises(errors.InvalidArgumentError): self.evaluate(math_ops.bincount([1, 2, 3, -1, 6, 8]))