diff options
Diffstat (limited to 'tensorflow/python/kernel_tests/cholesky_op_test.py')
-rw-r--r-- | tensorflow/python/kernel_tests/cholesky_op_test.py | 8 |
1 files changed, 4 insertions, 4 deletions
diff --git a/tensorflow/python/kernel_tests/cholesky_op_test.py b/tensorflow/python/kernel_tests/cholesky_op_test.py index 5369d2d5c4..d783522e82 100644 --- a/tensorflow/python/kernel_tests/cholesky_op_test.py +++ b/tensorflow/python/kernel_tests/cholesky_op_test.py @@ -311,7 +311,7 @@ class CholeskyBenchmark(test.Benchmark): if test.is_gpu_available(True): with ops.Graph().as_default(), \ session.Session() as sess, \ - ops.device("/gpu:0"): + ops.device("/device:GPU:0"): l = linalg_ops.cholesky(data) self.run_op_benchmark( sess, @@ -338,11 +338,11 @@ class CholeskyBenchmark(test.Benchmark): if test.is_gpu_available(True): _BenchmarkGrad( - MatrixInverseCompositeGrad, "composite_matrix_inverse", "/gpu:0") + MatrixInverseCompositeGrad, "composite_matrix_inverse", "/device:GPU:0") _BenchmarkGrad( - TriAngInvCompositeGrad, "composite_tri_ang_inverse", "/gpu:0") + TriAngInvCompositeGrad, "composite_tri_ang_inverse", "/device:GPU:0") _BenchmarkGrad( - TriAngSolveCompositeGrad, "composite_triangular_solve", "/gpu:0") + TriAngSolveCompositeGrad, "composite_triangular_solve", "/device:GPU:0") _BenchmarkGrad( MatrixInverseCompositeGrad, "composite_matrix_inverse", "/cpu:0") |