| 81 | |
| 82 | @test_util.run_gpu_only |
| 83 | def testGPU(self): |
| 84 | # tf.Tensor.as_gpu_device() moves a tensor to GPU. |
| 85 | x = constant_op.constant([[1., 2.], [3., 4.]]).gpu() |
| 86 | # Alternatively, tf.device() as a context manager places tensors and |
| 87 | # operations. |
| 88 | with ops.device('gpu:0'): |
| 89 | x += 1. |
| 90 | # Without a device context, heuristics are used to place ops. |
| 91 | # In this case, ops.reduce_mean runs on the GPU. |
| 92 | axis = range(x.shape.ndims) |
| 93 | m = math_ops.reduce_mean(x, axis) |
| 94 | # m is on GPU, bring it back to CPU and compare. |
| 95 | self.assertEqual(3.5, m.cpu().numpy()) |
| 96 | |
| 97 | def testListDevices(self): |
| 98 | # Expect at least one device. |