From 5608454c31bb298096bb6aa463b33baa2fa68f08 Mon Sep 17 00:00:00 2001 From: Ruoxin Sang Date: Thu, 4 Oct 2018 19:07:44 -0700 Subject: [PATCH] Add 'device' property to TPUMirroredVariable, so tf.train.init_from_checkpoint can be supported. PiperOrigin-RevId: 215843249 --- tensorflow/contrib/distribute/python/values.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/tensorflow/contrib/distribute/python/values.py b/tensorflow/contrib/distribute/python/values.py index 18ceba42c2a..0dd78ba185b 100644 --- a/tensorflow/contrib/distribute/python/values.py +++ b/tensorflow/contrib/distribute/python/values.py @@ -571,6 +571,10 @@ class TPUMirroredVariable(checkpointable.CheckpointableBase): ValueError("Device %s not found in %s (current device %s)" % (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() # function would normally see regular variables, not MirroredVariables. # However, the update function can still operate on wrapped MirroredVariables