From ae56c2ccc89b8677c655d0d2ec120ea27e9848c7 Mon Sep 17 00:00:00 2001
From: ceci3 <592712189@qq.com>
Date: Wed, 19 Jan 2022 12:56:57 +0000
Subject: [PATCH] fix save channel wise quant model

---
 .../paddle/fluid/contrib/slim/quantization/quantization_pass.py | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/python/paddle/fluid/contrib/slim/quantization/quantization_pass.py b/python/paddle/fluid/contrib/slim/quantization/quantization_pass.py
index 645feda21f0f3..efa000274d01a 100644
--- a/python/paddle/fluid/contrib/slim/quantization/quantization_pass.py
+++ b/python/paddle/fluid/contrib/slim/quantization/quantization_pass.py
@@ -1220,6 +1220,8 @@ def apply(self, graph):
             if op_node_desc.has_attr("quantization_type") and \
                 op_node_desc.attr("quantization_type") == "qat_with_weight":
                 if self._weight_quantize_type == 'channel_wise_abs_max':
+                    quant_axis = 1 if op_node.name() in \
+                        _channelwise_quant_axis1_ops else 0
                     self._insert_post_channel_dequant_op(graph, op_node,
                                                          quant_axis)
                 else: