forked from mindspore-Ecosystem/mindspore
Optimize performance of PyNative grad reduce
This commit is contained in:
parent
9605102dd1
commit
23898c735f
|
@ -65,9 +65,7 @@ def _tensors_allreduce(degree, mean, allgather, allreduce, allreduce_filter, gra
|
||||||
grad = allreduce(grad)
|
grad = allreduce(grad)
|
||||||
if mean:
|
if mean:
|
||||||
degree = F.scalar_cast(degree, F.dtype(grad))
|
degree = F.scalar_cast(degree, F.dtype(grad))
|
||||||
cast_op = P.Cast()
|
grad = F.tensor_mul(grad, F.cast(F.scalar_to_array(1.0 / degree), F.dtype(grad)))
|
||||||
mul_op = P.Mul()
|
|
||||||
grad = mul_op(grad, cast_op(F.scalar_to_array(1.0 / degree), F.dtype(grad)))
|
|
||||||
return grad
|
return grad
|
||||||
return grad
|
return grad
|
||||||
|
|
||||||
|
|
Loading…
Reference in New Issue