commit d0e21cd46860370f6146bf1ffdc07620744f2bbd
author Ken Franko <kfranko@google.com> Fri Mar 20 11:47:16 2020 -0700
committer TensorFlower Gardener <gardener@tensorflow.org> Fri Mar 20 11:58:13 2020 -0700
tree 76bc169a0d7782a93486a8180c98f17d14b9897d
parent 0ad3c881ff7724cf6275276670f2b7efe9dab382

Use correct variable _device attribute in Keras optimizer_v2.

PiperOrigin-RevId: 302069884
Change-Id: I32ff43f146c6f60d462d2713908c3cf258ace3de

tensorflow/python/keras/optimizer_v2/optimizer_v2.py

diff --git a/tensorflow/python/keras/optimizer_v2/optimizer_v2.py b/tensorflow/python/keras/optimizer_v2/optimizer_v2.py
index 2a4d4cf..d9f090c 100644
--- a/tensorflow/python/keras/optimizer_v2/optimizer_v2.py
+++ b/tensorflow/python/keras/optimizer_v2/optimizer_v2.py
@@ -699,8 +699,10 @@
   def _prepare(self, var_list):
     keys = set()
     for var in var_list:
-      var_devices = (getattr(var, "devices", None) or  # Distributed
-                     [var.device])                     # Regular
+      if isinstance(var, ds_values.DistributedValues):
+        var_devices = var._devices   # pylint: disable=protected-access
+      else:
+        var_devices = [var.device]
       var_dtype = var.dtype.base_dtype
       for var_device in var_devices:
         keys.add((var_device, var_dtype))