fix ActivationGradCompute on CPU

fixes apache#13333
larroy · Nov 20, 2018 · 66f44e4 · 66f44e4
1 parent 779bdc5
commit 66f44e4
Showing 1 changed file with 12 additions and 4 deletions.
diff --git a/src/operator/nn/activation-inl.h b/src/operator/nn/activation-inl.h
@@ -199,12 +199,20 @@ void ActivationGradCompute(const nnvm::NodeAttrs& attrs,
                            const std::vector<OpReqType>& req,
                            const std::vector<TBlob>& outputs) {
   const ActivationParam& param = nnvm::get<ActivationParam>(attrs.parsed);
+
+  const int act_type = param.act_type;
 #if (MXNET_USE_CUDNN == 1 || MXNET_USE_MKLDNN == 1)
-  bool relu = param.act_type == activation::kReLU;
-  CHECK_EQ(inputs.size(), relu ? 2U : 3U);
+  if (act_type != activation::kReLU && act_type != activation::kSoftSign) {
+    CHECK_EQ(inputs.size(), 3U);
+  } else {
+    CHECK_EQ(inputs.size(), 2U);
+  }
 #else
-  bool softsign = param.act_type == activation::kSoftSign;
-  CHECK_EQ(inputs.size(), softsign ? 3U : 2U);
+  if (act_type == activation::kSoftSign) {
+    CHECK_EQ(inputs.size(), 3U);
+  } else {
+    CHECK_EQ(inputs.size(), 2U);
+  }
 #endif
   CHECK_EQ(outputs.size(), 1U);
   CHECK_EQ(req.size(), 1U);