From 1506d266947b97c77ce9ada257bf9b811eee6dab Mon Sep 17 00:00:00 2001
From: Guanghua Yu <742925032@qq.com>
Date: Fri, 27 Aug 2021 13:49:18 +0800
Subject: [PATCH] Support the output of nn.Layer is dict type when in
 quantization training (#35180)

---
 python/paddle/nn/quant/quant_layers.py | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/python/paddle/nn/quant/quant_layers.py b/python/paddle/nn/quant/quant_layers.py
index 040b04f5e7..220e0c34a1 100644
--- a/python/paddle/nn/quant/quant_layers.py
+++ b/python/paddle/nn/quant/quant_layers.py
@@ -677,7 +677,8 @@ class MAOutputScaleLayer(layers.Layer):
     def forward(self, *inputs, **kwargs):
         out = self._layer(*inputs, **kwargs)
         # TODO (jc): support the ops of several outputs
-        if (isinstance(out, list) or isinstance(out, tuple)) and len(out) > 1:
+        if (isinstance(out, list) or isinstance(out, tuple) or
+                isinstance(out, dict)):
             return out
         else:
             return self._ma_output_scale(out)
-- 
GitLab