From 81e702ac835964227bbab0f4ac49ffddb6ff0a6d Mon Sep 17 00:00:00 2001
From: XGZhang <46363693+XGZhang11@users.noreply.github.com>
Date: Thu, 9 Sep 2021 19:16:14 +0800
Subject: [PATCH] quant: fix a export bug (#35410)

---
 .../fluid/contrib/slim/quantization/imperative/qat.py | 11 ++++++-----
 .../contrib/slim/quantization/quantization_pass.py    |  1 +
 2 files changed, 7 insertions(+), 5 deletions(-)

diff --git a/python/paddle/fluid/contrib/slim/quantization/imperative/qat.py b/python/paddle/fluid/contrib/slim/quantization/imperative/qat.py
index 790eff04c36..5620bdbb3a4 100644
--- a/python/paddle/fluid/contrib/slim/quantization/imperative/qat.py
+++ b/python/paddle/fluid/contrib/slim/quantization/imperative/qat.py
@@ -568,11 +568,12 @@ class ImperativeQuantizeOutputs(object):
                 out_scale = utils.fp_numpy_to_naive(out_scale)
 
                 if previous_op.type != "feed":
-                    argname, index = utils._get_output_name_index(previous_op,
-                                                                  in_var_name)
-                    previous_op._set_attr(argname + str(index) + "_threshold",
-                                          out_scale)
-                    previous_op._set_attr("out_threshold", out_scale)
+                    res = utils._get_output_name_index(previous_op, in_var_name)
+                    if res is not None:
+                        argname, index = res
+                        previous_op._set_attr(
+                            argname + str(index) + "_threshold", out_scale)
+                        previous_op._set_attr("out_threshold", out_scale)
 
                 for next_op in next_ops:
                     next_op._rename_input(out_var_name, in_var_name)
diff --git a/python/paddle/fluid/contrib/slim/quantization/quantization_pass.py b/python/paddle/fluid/contrib/slim/quantization/quantization_pass.py
index 01b54f8f13a..57a8abe69ec 100644
--- a/python/paddle/fluid/contrib/slim/quantization/quantization_pass.py
+++ b/python/paddle/fluid/contrib/slim/quantization/quantization_pass.py
@@ -134,6 +134,7 @@ _op_real_in_out_name = {
     "flatten": [["X"], ["Out"]],
     "flatten2": [["X"], ["Out"]],
     "unsqueeze2": [["X"], ["Out"]],
+    "flatten_contiguous_range": [['X'], ["Out", "XShape"]],
 }
 
 _conv_ops = ['conv2d', 'depthwise_conv2d', 'conv2d_transpose']
-- 
GitLab