fix expand_dims fall back when input's ndim is 0 (#16837)

apache · Nov 16, 2019 · 69fd018 · 69fd018
1 parent ec2f3a6
commit 69fd018
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/src/operator/nn/mkldnn/mkldnn_reshape.cc b/src/operator/nn/mkldnn/mkldnn_reshape.cc
@@ -127,7 +127,8 @@ void MKLDNNReshapeForward(const nnvm::NodeAttrs& attrs,
                           const NDArray &output) {
   // For mkldnn non-supported input, it shouldn't hold mkldnn memory, so let's simply fallback to
   // naive implement.
-  if (input.shape().ndim() > 4 || !SupportMKLDNNQuantize(input.dtype())) {
+  const int input_ndims = input.shape().ndim();
+  if ((input_ndims < 1 || input_ndims > 4) || !SupportMKLDNNQuantize(input.dtype())) {
     if (req != kWriteInplace) {
       FallBackCompute(UnaryOp::IdentityCompute<cpu>, attrs, ctx, {input}, {req}, {output});
     }