Add decodeOp validate

5 years ago · a8cc48c6b8
parent 42047764cf
commit a8cc48c6b8
1 changed files with 4 additions and 0 deletions
--- a/mindspore/ccsrc/minddata/dataset/kernels/image/decode_op.cc
+++ b/mindspore/ccsrc/minddata/dataset/kernels/image/decode_op.cc
@ -32,6 +32,10 @@ DecodeOp::DecodeOp(bool is_rgb_format) : is_rgb_format_(is_rgb_format) {

 Status DecodeOp::Compute(const std::shared_ptr<Tensor> &input, std::shared_ptr<Tensor> *output) {
  IO_CHECK(input, output);
+  // check the input tensor shape
+  if (input->Rank() != 1) {
+    RETURN_STATUS_UNEXPECTED("DecodeOp error: invalid input shape, only support 1D input.");
+  }
  if (is_rgb_format_) {  // RGB colour mode
    return Decode(input, output);
  } else {  // BGR colour mode