update

apache · Nov 26, 2018 · 76c1e1c · 76c1e1c
1 parent 21850b2
commit 76c1e1c
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 5 deletions.
diff --git a/src/operator/contrib/corner_pooling-inl.h b/src/operator/contrib/corner_pooling-inl.h
@@ -160,8 +160,11 @@ inline void corner_pool_grad(mshadow::Stream<cpu> *s,
                              const int corner_pooling_type,
                              OpReqType req_type,
                              DType *in_grad) {
-  mxnet_op::Kernel<mxnet_op::set_zero, cpu>::Launch(s, ishape.Size(),
-                                                    in_grad);
+  if (mxnet::kNullOp == req_type) return;
+  if (mxnet::kAddTo != req_type) {
+    mxnet_op::Kernel<mxnet_op::set_zero, cpu>::Launch(s, ishape.Size(),
+                                                      in_grad);
+  }
   const int height = ishape[2], width = ishape[3];
   const index_t data_offset = width * height;
   if (corner_pooling_type == 0 || corner_pooling_type == 1) {

diff --git a/src/operator/contrib/corner_pooling.cu b/src/operator/contrib/corner_pooling.cu
@@ -218,9 +218,11 @@ inline void corner_pool_grad(mshadow::Stream<gpu> *s,
                    const DType *out_data, const TShape &ishape,
                    const int corner_pooling_type, OpReqType req_type,
                    DType *in_grad) {
-  mxnet_op::Kernel<mxnet_op::set_zero, gpu>::Launch(s, ishape.Size(),
-                                                    in_grad);
-
+  if (mxnet::kNullOp == req_type) return;
+  if (mxnet::kAddTo != req_type) {
+    mxnet_op::Kernel<mxnet_op::set_zero, gpu>::Launch(s, ishape.Size(),
+                                                      in_grad);
+  }
   const int height = ishape[2], width = ishape[3];
   if (corner_pooling_type == 0 || corner_pooling_type == 1) {
     // top or bottom