未验证 提交 1506d266 编写于 作者: G Guanghua Yu 提交者: GitHub

Support the output of nn.Layer is dict type when in quantization training (#35180)

上级 50836e1d
...@@ -677,7 +677,8 @@ class MAOutputScaleLayer(layers.Layer): ...@@ -677,7 +677,8 @@ class MAOutputScaleLayer(layers.Layer):
def forward(self, *inputs, **kwargs): def forward(self, *inputs, **kwargs):
out = self._layer(*inputs, **kwargs) out = self._layer(*inputs, **kwargs)
# TODO (jc): support the ops of several outputs # 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 return out
else: else:
return self._ma_output_scale(out) return self._ma_output_scale(out)
......
Markdown is supported
0% .
You are about to add 0 people to the discussion. Proceed with caution.
先完成此消息的编辑!
想要评论请 注册