diff options
Diffstat (limited to 'tensorflow/tools/graph_transforms/quantize_nodes.cc')
-rw-r--r-- | tensorflow/tools/graph_transforms/quantize_nodes.cc | 2 |
1 files changed, 2 insertions, 0 deletions
diff --git a/tensorflow/tools/graph_transforms/quantize_nodes.cc b/tensorflow/tools/graph_transforms/quantize_nodes.cc index 2b85e7e83c..97e8f77616 100644 --- a/tensorflow/tools/graph_transforms/quantize_nodes.cc +++ b/tensorflow/tools/graph_transforms/quantize_nodes.cc @@ -759,6 +759,7 @@ Status QuantizeNodes(const GraphDef& input_graph_def, NodeDef reshape_dims; reshape_dims.set_op("Const"); reshape_dims.set_name(unique_input_name + "/reshape_dims"); + AddNodeInput("^" + input_name, &reshape_dims); SetNodeAttr("dtype", DT_INT32, &reshape_dims); Tensor reshape_dims_tensor(DT_INT32, {1}); reshape_dims_tensor.flat<int32>()(0) = -1; @@ -768,6 +769,7 @@ Status QuantizeNodes(const GraphDef& input_graph_def, NodeDef reduction_dims; reduction_dims.set_op("Const"); reduction_dims.set_name(unique_input_name + "/reduction_dims"); + AddNodeInput("^" + input_name, &reduction_dims); SetNodeAttr("dtype", DT_INT32, &reduction_dims); Tensor reduction_dims_tensor(DT_INT32, {1}); reduction_dims_tensor.flat<int32>()(0) = 0; |