add degree and mean to TrainOneStepCell

4 years ago · 0f844de1a9
parent aa9bee0ce3
commit 0f844de1a9
1 changed files with 3 additions and 3 deletions
--- 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