diff options
Diffstat (limited to 'tensorflow/python/kernel_tests/matrix_band_part_op_test.py')
| -rw-r--r-- | tensorflow/python/kernel_tests/matrix_band_part_op_test.py | 5 |
1 files changed, 3 insertions, 2 deletions
diff --git a/tensorflow/python/kernel_tests/matrix_band_part_op_test.py b/tensorflow/python/kernel_tests/matrix_band_part_op_test.py index 68d626de2c..a0ef3a607e 100644 --- a/tensorflow/python/kernel_tests/matrix_band_part_op_test.py +++ b/tensorflow/python/kernel_tests/matrix_band_part_op_test.py @@ -27,6 +27,7 @@ from tensorflow.python.ops import array_ops from tensorflow.python.ops import control_flow_ops from tensorflow.python.ops import gradient_checker from tensorflow.python.ops import variables +from tensorflow.python.platform import benchmark from tensorflow.python.platform import test as test_lib @@ -109,7 +110,7 @@ class MatrixBandPartBenchmark(test_lib.Benchmark): for shape_ in self.shapes: for limits in (-1, -1), (-1, 0), (0, -1), (2, 2): with ops.Graph().as_default(), \ - session.Session() as sess, \ + session.Session(config=benchmark.benchmark_config()) as sess, \ ops.device("/cpu:0"): matrix = variables.Variable(array_ops.ones(shape_)) band = array_ops.matrix_band_part(matrix, limits[0], limits[1]) @@ -123,7 +124,7 @@ class MatrixBandPartBenchmark(test_lib.Benchmark): if test_lib.is_gpu_available(True): with ops.Graph().as_default(), \ - session.Session() as sess, \ + session.Session(config=benchmark.benchmark_config()) as sess, \ ops.device("/gpu:0"): matrix = variables.Variable(array_ops.ones(shape_)) band = array_ops.matrix_band_part(matrix, limits[0], limits[1]) |