fix fill infer bug

4 years ago · c792d76c06
parent 27f7bfef50
commit c792d76c06
1 changed files with 3 additions and 1 deletions
--- a/mindspore/lite/nnacl/infer/fill_infer.c
+++ b/mindspore/lite/nnacl/infer/fill_infer.c
@ -38,7 +38,9 @@ int FillInferShape(const TensorC *const *inputs, size_t inputs_size, TensorC **o
  if (!parameter->infer_flag_) {
    return NNACL_INFER_INVALID;
  }
-
+  if (num_dims != 0 && dst_shape == NULL) {
+    return NNACL_INFER_INVALID;
+  }
  int output_shape[MAX_SHAPE_SIZE];
  size_t output_shape_size = 0;
  for (size_t i = 0; i < num_dims; i++) {