!13355 add degree and mean to TrainOneStepCell

From: @wangnan39
Reviewed-by: @kingxian,@zh_qh
Signed-off-by: @kingxian
pull/13355/MERGE
mindspore-ci-bot 4 years ago committed by Gitee
commit 4e43fb939a

@ -349,9 +349,9 @@ class TrainOneStepCell(Cell):
if self.parallel_mode in (ParallelMode.DATA_PARALLEL, ParallelMode.HYBRID_PARALLEL):
self.reducer_flag = True
if self.reducer_flag:
mean = _get_gradients_mean()
degree = _get_device_num()
self.grad_reducer = DistributedGradReducer(self.weights, mean, degree)
self.mean = _get_gradients_mean()
self.degree = _get_device_num()
self.grad_reducer = DistributedGradReducer(self.weights, self.mean, self.degree)
def construct(self, *inputs):
weights = self.weights

Loading…
Cancel
Save