Fix error when getting optimizer variables with distribution strategy
- add `_in_graph_mode` property to DistributedVariable
PiperOrigin-RevId: 210177702
diff --git a/tensorflow/contrib/distribute/python/examples/keras_mnist.py b/tensorflow/contrib/distribute/python/examples/keras_mnist.py
index e33e6fd..a20069c 100644
--- a/tensorflow/contrib/distribute/python/examples/keras_mnist.py
+++ b/tensorflow/contrib/distribute/python/examples/keras_mnist.py
@@ -109,8 +109,7 @@
# `distribute` argument. `fit`, `evaluate` and `predict` will be distributed
# based on the strategy instantiated.
model.compile(loss=tf.keras.losses.categorical_crossentropy,
- optimizer=tf.train.GradientDescentOptimizer(
- learning_rate=0.001),
+ optimizer=tf.train.RMSPropOptimizer(learning_rate=0.001),
metrics=['accuracy'],
distribute=strategy)
diff --git a/tensorflow/contrib/distribute/python/values.py b/tensorflow/contrib/distribute/python/values.py
index e73d9c1..3ccaa26 100644
--- a/tensorflow/contrib/distribute/python/values.py
+++ b/tensorflow/contrib/distribute/python/values.py
@@ -304,6 +304,10 @@
self._primary_var.op.type)
return self.get().op
+ @property
+ def _in_graph_mode(self):
+ return self._primary_var._in_graph_mode # pylint: disable=protected-access
+
def read_value(self):
return distribution_strategy_context.get_distribution_strategy().read_var(
self)