From dc1e7b8a635426db809a094b313cc6be40ab9276 Mon Sep 17 00:00:00 2001 From: Akshay Modi Date: Fri, 9 Feb 2018 14:36:17 -0800 Subject: [PATCH] Fix read_variable_op GPU test PiperOrigin-RevId: 185194768 --- tensorflow/python/eager/benchmarks_test.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tensorflow/python/eager/benchmarks_test.py b/tensorflow/python/eager/benchmarks_test.py index 9aceaba..b56cbe8 100644 --- a/tensorflow/python/eager/benchmarks_test.py +++ b/tensorflow/python/eager/benchmarks_test.py @@ -425,7 +425,7 @@ class MicroBenchmarks(test.Benchmark): if not context.num_gpus(): return with context.device(GPU): - m = resource_variable_ops.ResourceVariable(self._m_2_by_2) + m = resource_variable_ops.ResourceVariable(self._m_2_by_2.gpu()) self._benchmark_read_variable(m, num_iters=self._num_iters_2_by_2) def benchmark_read_variable_op_with_tape_2_by_2_CPU(self): @@ -438,7 +438,7 @@ class MicroBenchmarks(test.Benchmark): if not context.num_gpus(): return with context.device(GPU): - m = resource_variable_ops.ResourceVariable(self._m_2_by_2) + m = resource_variable_ops.ResourceVariable(self._m_2_by_2.gpu()) self._benchmark_read_variable_with_tape( m, num_iters=self._num_iters_2_by_2) -- 2.7.4