[Op] Prevent inconsistent number of Ops and devices during distribute…

…d training.

Signed-off-by: JunqiHu <[email protected]>

tensorflow/python/training/saver.py

@@ -550,8 +550,14 @@ def _GroupByDevices(self, saveables):
     """
     per_device = collections.defaultdict(lambda: [])
     for saveable in saveables:
-      canonical_device = set(
-          pydev.canonical_name(spec.tensor.device) for spec in saveable.specs)
+      canonical_device = set()
+      for spec in saveable.specs:
+        device_name = pydev.canonical_name(spec.tensor.device)
+        device_idx = device_name.find("/device")
+        if device_idx != -1:
+          canonical_device.add(device_name[:device_idx])
+        else:
+          canonical_device.add(device_name)
       if len(canonical_device) != 1:
         raise ValueError("All tensors of a saveable object must be "
                          "on the same device: %s" % saveable.name)