inplace sum.

apache · Mar 29, 2018 · 8cce42e · 8cce42e
1 parent c8db046
commit 8cce42e
Showing 1 changed file with 2 additions and 5 deletions.
diff --git a/src/operator/nn/mkldnn/mkldnn_base.cc b/src/operator/nn/mkldnn/mkldnn_base.cc
@@ -118,11 +118,8 @@ void CommitOutput(const NDArray &arr, const mkldnn_output_t &res) {
   } else if (res.first == AddBack) {
     auto mem = arr.GetMKLDNNData(res.second->get_primitive_desc());
     CHECK(mem != nullptr);
-    // We have to allocate new memory for the sum result.
-    auto sum_res = TmpMemMgr::Get()->Alloc(
-        res.second->get_primitive_desc());
-    op::Sum(*res.second, *mem, *sum_res);
-    const_cast<NDArray &>(arr).CopyFrom(*sum_res);
+    CHECK(mem->get_primitive_desc() == res.second->get_primitive_desc());
+    op::Sum(*res.second, *mem, *mem);
   }
 }