aboutsummaryrefslogtreecommitdiffhomepage
path: root/tensorflow/python/kernel_tests/cond_v2_test.py
diff options
context:
space:
mode:
authorGravatar A. Unique TensorFlower <gardener@tensorflow.org>2018-09-13 01:35:35 -0700
committerGravatar TensorFlower Gardener <gardener@tensorflow.org>2018-09-13 01:38:41 -0700
commit567de999ae29a2cfb30132f82178006fe5688d6b (patch)
treeb5e934afc4ad03810cb2f956f72debebddb341cb /tensorflow/python/kernel_tests/cond_v2_test.py
parent7f3938deb393f7688cd364b630afdd9338460299 (diff)
Change test to use 2 CPU devices instead of GPU.
General cleanup: testDeviceInAndOutOfCond uses a GPU in a CPU only test build resulting in all operations run on the same device even though the graph is for multiple devices. PiperOrigin-RevId: 212775360
Diffstat (limited to 'tensorflow/python/kernel_tests/cond_v2_test.py')
-rw-r--r--tensorflow/python/kernel_tests/cond_v2_test.py8
1 files changed, 5 insertions, 3 deletions
diff --git a/tensorflow/python/kernel_tests/cond_v2_test.py b/tensorflow/python/kernel_tests/cond_v2_test.py
index 18a1b230a0..a1efecf28a 100644
--- a/tensorflow/python/kernel_tests/cond_v2_test.py
+++ b/tensorflow/python/kernel_tests/cond_v2_test.py
@@ -892,11 +892,13 @@ class CondV2ColocationGroupAndDeviceTest(test.TestCase):
def testDeviceInAndOutOfCond(self):
with ops.Graph().as_default() as g:
- with self.test_session(graph=g):
+ with self.test_session(
+ graph=g, config=config_pb2.ConfigProto(device_count={"CPU": 2})):
+
def fn2():
- with ops.device("/device:GPU:0"):
+ with ops.device("/device:CPU:1"):
c = constant_op.constant(3.0)
- self.assertEqual("/device:GPU:0", c.op.device)
+ self.assertEqual("/device:CPU:1", c.op.device)
return c
with ops.device("/device:CPU:0"):