From 0f844de1a9a9b351ad41a795ac9b29d11fb868a5 Mon Sep 17 00:00:00 2001 From: "wangnan39@huawei.com" Date: Mon, 15 Mar 2021 21:28:46 +0800 Subject: [PATCH] add degree and mean to TrainOneStepCell --- mindspore/nn/wrap/cell_wrapper.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/mindspore/nn/wrap/cell_wrapper.py b/mindspore/nn/wrap/cell_wrapper.py index f0d822fa3d..a8cde8801e 100644 --- a/mindspore/nn/wrap/cell_wrapper.py +++ b/mindspore/nn/wrap/cell_wrapper.py @@ -339,9 +339,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