diff --git a/mindspore/ccsrc/backend/kernel_compiler/cpu/apply_adagrad_cpu_kernel.cc b/mindspore/ccsrc/backend/kernel_compiler/cpu/apply_adagrad_cpu_kernel.cc index 34d4c1ee16..faa08392cd 100644 --- a/mindspore/ccsrc/backend/kernel_compiler/cpu/apply_adagrad_cpu_kernel.cc +++ b/mindspore/ccsrc/backend/kernel_compiler/cpu/apply_adagrad_cpu_kernel.cc @@ -38,9 +38,9 @@ bool ApplyAdagradCPUKernel::Launch(const std::vector &inputs, const CheckParam(inputs, outputs); if (dtype_ == kNumberTypeFloat16) { - LaunchKernel(inputs); + LaunchKernel(inputs, outputs); } else if (dtype_ == kNumberTypeFloat32) { - LaunchKernel(inputs); + LaunchKernel(inputs, outputs); } return true; @@ -67,7 +67,8 @@ void ApplyAdagradCPUKernel::CheckParam(const std::vector &inputs, co } template -void ApplyAdagradCPUKernel::LaunchKernel(const std::vector &inputs) { +void ApplyAdagradCPUKernel::LaunchKernel(const std::vector &inputs, + const std::vector &outputs) { auto var = reinterpret_cast(inputs[0]->addr); auto accum = reinterpret_cast(inputs[1]->addr); auto lr = reinterpret_cast(inputs[2]->addr); @@ -96,6 +97,17 @@ void ApplyAdagradCPUKernel::LaunchKernel(const std::vector &inputs) for (auto &it : threads) { it.join(); } + + // Copy result to output tensor + auto output_var = reinterpret_cast(outputs[0]->addr); + auto output_accum = reinterpret_cast(outputs[1]->addr); + if (memcpy_s(output_var, outputs[0]->size, var, inputs[0]->size) != EOK) { + MS_LOG(EXCEPTION) << "Launch kernel error: memcpy failed."; + } + + if (memcpy_s(output_accum, outputs[1]->size, accum, inputs[1]->size) != EOK) { + MS_LOG(EXCEPTION) << "Launch kernel error: memcpy failed."; + } } template diff --git a/mindspore/ccsrc/backend/kernel_compiler/cpu/apply_adagrad_cpu_kernel.h b/mindspore/ccsrc/backend/kernel_compiler/cpu/apply_adagrad_cpu_kernel.h index f5ebe37fbf..4af4a982b2 100644 --- a/mindspore/ccsrc/backend/kernel_compiler/cpu/apply_adagrad_cpu_kernel.h +++ b/mindspore/ccsrc/backend/kernel_compiler/cpu/apply_adagrad_cpu_kernel.h @@ -36,7 +36,7 @@ class ApplyAdagradCPUKernel : public CPUKernel { private: static void CheckParam(const std::vector &inputs, const std::vector &outputs); template - void LaunchKernel(const std::vector &inputs); + void LaunchKernel(const std::vector &inputs, const std::vector &outputs); template void LaunchApplyAdagrad(T const var, T const accum, const T lr, const T gradient, size_t start, size_t end); bool update_slots_{true}; diff --git a/tests/st/ops/cpu/test_apply_adagrad_op.py b/tests/st/ops/cpu/test_apply_adagrad_op.py index ff7f284cf3..dde952547b 100644 --- a/tests/st/ops/cpu/test_apply_adagrad_op.py +++ b/tests/st/ops/cpu/test_apply_adagrad_op.py @@ -36,8 +36,7 @@ class Net(nn.Cell): self.accum = Parameter(Tensor(accum_np), name="accum") def construct(self, lr, grad): - self.apply_adagrad(self.var, self.accum, lr, grad) - return self.var, self.accum + return self.apply_adagrad(self.var, self.accum, lr, grad) @pytest.mark.level0