fix reduce_mean prim float16 bug (#54836)

7f4cdd55 · cyber-pioneer · GitHub · a35aa8cf · 7f4cdd55
隐藏空白更改
内联并排

Showing with 1 addition and 1 deletion

python/paddle/incubate/autograd/composite_rules.py python/paddle/incubate/autograd/composite_rules.py +1 -1

未找到文件。
--- a/python/paddle/incubate/autograd/composite_rules.py
+++ b/python/paddle/incubate/autograd/composite_rules.py
@@ -248,7 +248,7 @@ def mean_composite(x, axis, keepdim):
    from paddle.fluid.data_feeder import convert_dtype

    dtype = convert_dtype(x.dtype)
-    if dtype == ["float16", "uint16"]:
+    if dtype in ["float16", "uint16"]:
        is_amp = True
        x = cast(x, "float32")