Make Saver work with GPU variables in Eager mode.

PiperOrigin-RevId: 168714847

1 files changed, 10 insertions, 1 deletions

diff --git a/tensorflow/python/training/saver.py b/tensorflow/python/training/saver.py
index dc9bd3a8c2..72ec106d2b 100644
--- a/tensorflow/python/training/saver.py
+++ b/tensorflow/python/training/saver.py
@@ -167,8 +167,17 @@ class BaseSaverBuilder(object):
         self.handle_op = var.op.inputs[0]
         tensor = var
       elif isinstance(var, resource_variable_ops.ResourceVariable):
+
+        def _read_variable_closure(v):
+          def f():
+            with ops.device(v.device):
+              x = v.read_value()
+            with ops.device("/device:CPU:0"):
+              return array_ops.identity(x)
+          return f
+
         self.handle_op = var.handle
-        tensor = var.read_value
+        tensor = _read_variable_closure(var)
       else:
         raise ValueError(
             "Saveable is neither a resource variable nor a read operation."