From 1dffb00347a80d41c91eb43d0416d2d4ee0d9bc1 Mon Sep 17 00:00:00 2001 From: Akshay Modi Date: Wed, 10 Jul 2019 11:32:24 -0700 Subject: [PATCH] Fix the device placement test - check the variable's device, instead of the read tensor's device. PiperOrigin-RevId: 257447029 --- tensorflow/python/eager/def_function_test.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/tensorflow/python/eager/def_function_test.py b/tensorflow/python/eager/def_function_test.py index 46f45395272..4a7d6fe4e9e 100644 --- a/tensorflow/python/eager/def_function_test.py +++ b/tensorflow/python/eager/def_function_test.py @@ -580,7 +580,7 @@ class DefFunctionTest(test.TestCase): # TODO(b/137148281): reenable @test_util.run_gpu_only - def DISABLED_testDeviceAnnotationRespected(self): + def testDeviceAnnotationRespected(self): a = [] @def_function.function() @@ -590,13 +590,13 @@ class DefFunctionTest(test.TestCase): (2, 2), maxval=1000000, dtype=dtypes.int64) if not a: - with ops.device("CPU:0"): + with ops.device('CPU:0'): a.append(resource_variable_ops.ResourceVariable(initial_value)) return a[0].read_value() created_variable_read = create_variable() - self.assertRegexpMatches(created_variable_read.device, "CPU") + self.assertRegexpMatches(a[0].device, 'CPU') def testDecorate(self): func = def_function.function(lambda: 1)