diff --git a/src/operator/tensor/elemwise_sum.h b/src/operator/tensor/elemwise_sum.h index e89e9d799903..259c80ddddac 100644 --- a/src/operator/tensor/elemwise_sum.h +++ b/src/operator/tensor/elemwise_sum.h @@ -94,7 +94,7 @@ void ElementWiseSumCompute_(const nnvm::NodeAttrs& attrs, Kernel::Launch(s, out_size, out_dptr, req[0], in_0_dptr); for (size_t i = 1; i < size; ++i) { DType* in_dptr = in_data[i].dptr(); - Kernel::Launch(s, out_size, out_dptr, req[0], out_dptr, in_dptr); + Kernel::Launch(s, out_size, out_dptr, kWriteTo, out_dptr, in_dptr); } break; } diff --git a/tests/python/unittest/test_operator.py b/tests/python/unittest/test_operator.py index 6cbbc5dd0509..df4a77f66ea4 100644 --- a/tests/python/unittest/test_operator.py +++ b/tests/python/unittest/test_operator.py @@ -9875,6 +9875,25 @@ def test_grad(input_shape, kernel, stride=1, dilate=1, pad=0): pad = 1 ) +def test_elemwise_sum_for_gradient_accumulation(): + for nrepeat in range(1, 10): + stored_grad = dict() + for grad_req in ['write', 'add']: + a = mx.nd.array([1]) + b = mx.nd.array([2]) + if grad_req == 'write': + a.attach_grad(grad_req='write') + elif grad_req == 'add': + a.attach_grad(grad_req='add') + a.grad[:] = 0 + with mx.autograd.record(): + for _ in range(nrepeat): + b = b * a + b.backward() + stored_grad[grad_req] = a.grad.asscalar() + assert stored_grad['write'] == stored_grad['add'] + assert stored_grad['write'] == 2 * nrepeat + if __name__ == '__main__': import nose