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

Merge pull request #41357 from kenfranko/cherrypicks_4ECL0

Browse Source 
2.3-rc2 cherry-pick request: Correctly set the experimental_io_device when restoring variable from a checkpoint.
This commit is contained in:
Mihai Maruseac 2020-07-17 18:48:31 +00:00 committed by GitHub
commit 8edf01fbdc
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 3 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

5
tensorflow/python/training/tracking/base.py
View File

@ -293,9 +293,10 @@ class CheckpointPosition(object):
checkpoint_key = serialized_tensor.checkpoint_key
dtype = self._checkpoint.dtype_map[checkpoint_key]
base_type = dtype.base_dtype
io_device = self._checkpoint.options.experimental_io_device or "cpu:0"
with ops.init_scope():
with ops.device("/cpu:0"):
# Run the restore itself on the CPU.
with ops.device(io_device):
# Run the restore itself on the io_device(CPU or specified).
value, = io_ops.restore_v2(
prefix=self._checkpoint.save_path_tensor,
tensor_names=[checkpoint_key],