experiments/STT-tensorflow
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

Keep the mapping of distributed variables' values in map_resources().

Browse Source 
PiperOrigin-RevId: 279422939
Change-Id: Ie22c960bfb87f3dd8aabb14cb788d1a0a936efd1
This commit is contained in:
Rick Chao 2019-11-08 17:21:36 -08:00 committed by TensorFlower Gardener
parent ecd935d643
commit c82df39b0c
1 changed files with 4 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
tensorflow/python/saved_model/save.py
View File

@ -260,6 +260,10 @@ class _SaveableView(object):
obj_to_copy)
if ds_values.is_distributed_variable(obj):
self.captured_tensor_node_ids[obj] = node_id
for v in obj.values:
object_map[v] = new_variable
resource_map[v.handle] = new_variable.handle
self.captured_tensor_node_ids[v.handle] = node_id
object_map[obj] = new_variable
resource_map[obj.handle] = new_variable.handle
self.captured_tensor_node_ids[obj.handle] = node_id