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

Add 'device' property to TPUMirroredVariable, so tf.train.init_from_checkpoint can be supported.

Browse Source 
PiperOrigin-RevId: 215843249
This commit is contained in:
Ruoxin Sang 2018-10-04 19:07:44 -07:00 committed by TensorFlower Gardener
parent 83ff640fa5
commit 5608454c31
1 changed files with 4 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
tensorflow/contrib/distribute/python/values.py
View File

@ -571,6 +571,10 @@ class TPUMirroredVariable(checkpointable.CheckpointableBase):
ValueError("Device %s not found in %s (current device %s)" % ValueError("Device %s not found in %s (current device %s)" %
(device, self._index.keys(), device_util.current())), e) (device, self._index.keys(), device_util.current())), e)
@property
def device(self):
return self._get().device
# The arguments to update() are automatically unwrapped so the update() # The arguments to update() are automatically unwrapped so the update()
# function would normally see regular variables, not MirroredVariables. # function would normally see regular variables, not MirroredVariables.
# However, the update function can still operate on wrapped MirroredVariables # However, the update function can still operate on wrapped MirroredVariables