From da10fb3bad70d5c4cdb6f7affa7019c82f3ae593 Mon Sep 17 00:00:00 2001 From: ceci3 Date: Tue, 20 Sep 2022 15:57:44 +0800 Subject: [PATCH] update ma output link in graph (#46142) --- .../fluid/contrib/slim/quantization/quantization_pass.py | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/python/paddle/fluid/contrib/slim/quantization/quantization_pass.py b/python/paddle/fluid/contrib/slim/quantization/quantization_pass.py index 520c094798..4ba7a164c3 100644 --- a/python/paddle/fluid/contrib/slim/quantization/quantization_pass.py +++ b/python/paddle/fluid/contrib/slim/quantization/quantization_pass.py @@ -1531,8 +1531,16 @@ class OutScaleForTrainingPass(object): attrs=attrs, inputs=ins, outputs=outs) + + next_op_node = None + if len(in_node.outputs) > 0: + next_op_node = in_node.outputs[0] + graph.link_to(in_node, scale_op_node) graph.link_to(scale_op_node, scale_node) + if next_op_node: + graph.link_to(scale_node, next_op_node) + if not self._is_test: graph.link_to(state_in_node, scale_op_node) graph.link_to(accum_in_node, scale_op_node) -- GitLab