fix bug

apache · May 13, 2019 · 07dd00b · 07dd00b
1 parent 200e8bf
commit 07dd00b
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/src/operator/nn/layer_norm.cu b/src/operator/nn/layer_norm.cu
@@ -229,8 +229,8 @@ void LayerNormGPUContig(const LayerNormParam param,
                         const std::vector<TBlob>& outputs) {
   using namespace mshadow;
   CHECK_EQ(inputs.size(), 3U);
-  mxnet::TShape data_shape(2);
-  mxnet::TShape mean_shape(1);
+  mxnet::TShape data_shape(2, 0);
+  mxnet::TShape mean_shape(1, 0);
   size_t in_ndim = inputs[layernorm::kData].ndim();
   data_shape[0] = mean_shape[0] = inputs[layernorm::kData].shape_.ProdShape(0, in_ndim - 1);
   data_shape[1] = inputs[layernorm::kData].shape_[in_ndim - 1];