From 57e24820893e33bd2021460a037b87c093b0e50e Mon Sep 17 00:00:00 2001 From: Hector Li Date: Wed, 31 Jul 2019 21:21:49 -0700 Subject: [PATCH] Fix a bug in Expand cuda op implementation. (#1528) Description: crash if the output shape has 0 in it. because the code to / output_shape[i] Fix: If the output shape has 0 which means output_shape.Size() is 0, so output should be null. --- onnxruntime/core/providers/cuda/tensor/expand.cc | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/onnxruntime/core/providers/cuda/tensor/expand.cc b/onnxruntime/core/providers/cuda/tensor/expand.cc index 599fa39f406b2..897e981877ad0 100644 --- a/onnxruntime/core/providers/cuda/tensor/expand.cc +++ b/onnxruntime/core/providers/cuda/tensor/expand.cc @@ -21,6 +21,11 @@ Status Expand::ComputeInternal(OpKernelContext* ctx) const { ORT_RETURN_IF_ERROR(ComputeOutputShape(Node().Name(), input0.Shape(), output_dims, output_shape)); auto rank = output_shape.NumDimensions(); auto& output_tensor = *ctx->Output(0, output_shape); + + if (0 == output_shape.Size()) { + return Status::OK(); + } + auto input_shape = input0.Shape().GetDims(); // pad input_dims with 1 to make ranks match @@ -40,6 +45,8 @@ Status Expand::ComputeInternal(OpKernelContext* ctx) const { for (auto i = 0; i < rank; i++) { in_span[i] = fast_divmod(static_cast(input_shape[i])); out_span[i] = fast_divmod(static_cast(output_shape[i])); + // output_shape[i] won't be 0 here, it's covered in (0 == output_shape.Size()) + // a null output will be returned for that case subdim_size /= output_shape[i]; sdm_span[i] = static_cast(subdim_size); }