fix the if condition for LayerNorm (apache#15094)

haohuanw · Jun 23, 2019 · 7b4dd4b · 7b4dd4b
1 parent 32f176e
commit 7b4dd4b
Showing 1 changed file with 3 additions and 2 deletions.
diff --git a/src/operator/nn/layer_norm.cc b/src/operator/nn/layer_norm.cc
@@ -86,8 +86,9 @@ void LayerNormComputeMKL(const nnvm::NodeAttrs& attrs,
   CHECK_EQ(inputs.size(), 3U);
   int axis = GetRealAxis(param.axis, inputs[0].ndim());
 
-  if (axis == (inputs[layernorm::kData].ndim() - 1) ||
-      (inputs[0].type_flag_ != kFloat32 && inputs[0].type_flag_ != kFloat64)) {
+  // This optimization only applys for LayerNorm on the last dimension with dtype FP32 or FP64.
+  if (axis == (inputs[layernorm::kData].ndim() - 1) &&
+      (inputs[0].type_flag_ == kFloat32 || inputs[0].type_flag_ == kFloat64)) {
     // Compute necessary data for the reduce operation.
     mxnet::TShape red_src_shape, red_dst_shape;
     BroadcastReduceShapeCompact(inputs[layernorm::kData].shape_, outputs[layernorm::kMean].shape_,