experiments/STT-tensorflow
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

Update saver.py

Browse Source 
Fix device placement of save_op for ResourceVariable.
This commit is contained in:
Jacker 2018-07-20 10:09:16 +08:00 committed by GitHub
parent 6155633318
commit 0c11bcb5f3
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 4 additions and 7 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

11
tensorflow/python/training/saver.py
View File

@ -126,13 +126,10 @@ class BaseSaverBuilder(object):
def f():
with ops.device(v.device):
x = v.read_value()
# To allow variables placed on non-CPU devices to be checkpointed,
# we copy them to CPU on the same machine first.
device_spec = pydev.DeviceSpec().parse_from_string(v.device)
device_spec.merge_from(
pydev.DeviceSpec().parse_from_string("/device:CPU:0"))
with ops.device(device_spec.to_string()):
return array_ops.identity(x)
# To allow variables placed on non-CPU devices to be checkpointed,
# we copy them to CPU on the same machine first.
with ops.device("/device:CPU:0"):
return array_ops.identity(x)
return f
self.handle_op = var.handle