fix save channel wise quant model (PaddlePaddle#39054)

veyron95 · Jan 21, 2022 · ab1abd4 · ab1abd4
1 parent 064bc4b
commit ab1abd4
Showing 1 changed file with 2 additions and 0 deletions.
diff --git 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: