Internal code refactoring
@@ -377,6 +377,17 @@ class AutoCastVariable(tf.Variable, tf.__internal__.types.Tensor):
obj_map[self] = obj_map[self._variable]
return obj_map, resource_map
+ def _export_to_saved_model_graph(
+ self, object_map, tensor_map, options, **kwargs
+ ):
+ # By delegating this method to the wrapped variable, SavedModel with
+ # AutoCastVariables are identical to SavedModel with normal variables.
+ resource_list = self._variable._export_to_saved_model_graph(
+ object_map, tensor_map, options, **kwargs
+ )
+ object_map[self] = object_map[self._variable]
+ return resource_list
+
# TODO(reedwm): Maybe encode the fact the variable is an AutoCastVariable in
# to_proto().
def to_proto(self, export_scope=None):
