From 27935b664b5a2f64353406c824ae6d2561c08898 Mon Sep 17 00:00:00 2001 From: Rick Chao Date: Wed, 23 Oct 2019 17:13:28 -0700 Subject: [PATCH] In resource_variable_ops.copy_to_graph_uninitialized, instantiate `UninitializedVariable` under var.device's scope. It was discovered that when using tf.saved_model.save with MultiWorkerMirroredStrategy, it is possible that a variable does not have device attribute and tf.distribute code errors out. This fixes it and a test that would have failed is added to cover such use case. PiperOrigin-RevId: 276383534 Change-Id: I467bdaf5946da7dd1f3722c91a586830566608de --- tensorflow/python/saved_model/save.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/tensorflow/python/saved_model/save.py b/tensorflow/python/saved_model/save.py index 2ee97c445c5..0672f223fab 100644 --- a/tensorflow/python/saved_model/save.py +++ b/tensorflow/python/saved_model/save.py @@ -259,7 +259,9 @@ class _SaveableView(object): # created component variables. new_vars = [] for v in obj.values: - new_variable = resource_variable_ops.copy_to_graph_uninitialized(v) + # Ensure the variables are created with device attribute set. + with ops.device(v.device): + new_variable = resource_variable_ops.copy_to_graph_uninitialized(v) object_map[v] = new_variable new_vars.append(new_variable) resource_map[v.handle] = new_variable.handle