diff options
Diffstat (limited to 'tensorflow/python/kernel_tests/bucketize_op_test.py')
| -rw-r--r-- | tensorflow/python/kernel_tests/bucketize_op_test.py | 8 |
1 files changed, 4 insertions, 4 deletions
diff --git a/tensorflow/python/kernel_tests/bucketize_op_test.py b/tensorflow/python/kernel_tests/bucketize_op_test.py index 6db3592055..e612b1c134 100644 --- a/tensorflow/python/kernel_tests/bucketize_op_test.py +++ b/tensorflow/python/kernel_tests/bucketize_op_test.py @@ -31,7 +31,7 @@ class BucketizationOpTest(test.TestCase): constant_op.constant([-5, 0, 2, 3, 5, 8, 10, 11, 12]), boundaries=[0, 3, 8, 11]) expected_out = [0, 1, 1, 2, 2, 3, 3, 4, 4] - with self.test_session() as sess: + with self.test_session(use_gpu=True) as sess: self.assertAllEqual(expected_out, sess.run(op)) def testFloat(self): @@ -39,7 +39,7 @@ class BucketizationOpTest(test.TestCase): constant_op.constant([-5., 0., 2., 3., 5., 8., 10., 11., 12.]), boundaries=[0., 3., 8., 11.]) expected_out = [0, 1, 1, 2, 2, 3, 3, 4, 4] - with self.test_session() as sess: + with self.test_session(use_gpu=True) as sess: self.assertAllEqual(expected_out, sess.run(op)) def test2DInput(self): @@ -47,13 +47,13 @@ class BucketizationOpTest(test.TestCase): constant_op.constant([[-5, 0, 2, 3, 5], [8, 10, 11, 12, 0]]), boundaries=[0, 3, 8, 11]) expected_out = [[0, 1, 1, 2, 2], [3, 3, 4, 4, 1]] - with self.test_session() as sess: + with self.test_session(use_gpu=True) as sess: self.assertAllEqual(expected_out, sess.run(op)) def testInvalidBoundariesOrder(self): op = math_ops._bucketize( constant_op.constant([-5, 0]), boundaries=[0, 8, 3, 11]) - with self.test_session() as sess: + with self.test_session(use_gpu=True) as sess: with self.assertRaisesRegexp( errors_impl.InvalidArgumentError, "Expected sorted boundaries"): sess.run(op) |