Optimize performance of PyNative grad reduce

4 years ago · 23898c735f
parent 9605102dd1
commit 23898c735f
1 changed files with 1 additions and 3 deletions
--- a/mindspore/nn/wrap/grad_reducer.py
+++ b/mindspore/nn/wrap/grad_reducer.py
@ -65,9 +65,7 @@ def _tensors_allreduce(degree, mean, allgather, allreduce, allreduce_filter, gra
        grad = allreduce(grad)
        if mean:
            degree = F.scalar_cast(degree, F.dtype(grad))
-            cast_op = P.Cast()
-            mul_op = P.Mul()
-            grad = mul_op(grad, cast_op(F.scalar_to_array(1.0 / degree), F.dtype(grad)))
+            grad = F.tensor_mul(grad, F.cast(F.scalar_to_array(1.0 / degree), F.dtype(grad)))
        return grad
    return grad