diff --git a/onnxruntime/contrib_ops/cpu/image_scaler.h b/onnxruntime/contrib_ops/cpu/image_scaler.h index 865bca51f1e85..12e47670a8f62 100644 --- a/onnxruntime/contrib_ops/cpu/image_scaler.h +++ b/onnxruntime/contrib_ops/cpu/image_scaler.h @@ -36,7 +36,11 @@ class ImageScaler final : public OpKernel { const int64_t H = dims[2]; const int64_t W = dims[3]; - if (!bias_.empty() && bias_.size() != static_cast(C)) { + // The loop below reads bias_[nc % C] for every channel, so the bias must have exactly one entry + // per channel. An empty bias is not a valid "no bias" state here: the constructor already fails + // when the attribute is absent, and a present-but-empty attribute would otherwise index an + // empty vector. + if (bias_.size() != static_cast(C)) { return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "Bias size (", bias_.size(), ") does not match the number of channels (", C, ")"); } diff --git a/onnxruntime/contrib_ops/cuda/tensor/image_scaler.cc b/onnxruntime/contrib_ops/cuda/tensor/image_scaler.cc index befad5661c43f..8c0b60bedb6d5 100644 --- a/onnxruntime/contrib_ops/cuda/tensor/image_scaler.cc +++ b/onnxruntime/contrib_ops/cuda/tensor/image_scaler.cc @@ -47,7 +47,9 @@ Status ImageScaler::ComputeInternal(OpKernelContext* context) const { const int64_t C = dims[1]; // dims are NCHW - if (!bias_.empty() && bias_.size() != static_cast(C)) { + // The kernel indexes bias_data[c] for every channel, so the bias must have exactly one entry per + // channel. An empty bias would leave b_data_ pointing at a zero-sized allocation. + if (bias_.size() != static_cast(C)) { return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "Bias size (", bias_.size(), ") does not match the number of channels (", C, ")"); } diff --git a/onnxruntime/test/contrib_ops/tensor_op_test.cc b/onnxruntime/test/contrib_ops/tensor_op_test.cc index 2922e6517943e..d3d2b4a7d330a 100644 --- a/onnxruntime/test/contrib_ops/tensor_op_test.cc +++ b/onnxruntime/test/contrib_ops/tensor_op_test.cc @@ -98,6 +98,31 @@ TEST(ImageScalerContribOpTest, ImageScalerTest) { test.Run(); } +TEST(ImageScalerContribOpTest, ImageScalerEmptyBias) { + if (DefaultDmlExecutionProvider().get() != nullptr) { + GTEST_SKIP() << "Skipping because of the following error: AbiCustomRegistry.cpp(507): The parameter is incorrect."; + } + + constexpr int64_t N = 1, C = 2, H = 2, W = 2; + std::vector X = { + 1.0f, 3.0f, + 3.0f, 5.0f, + + 3.0f, 5.0f, + 7.0f, 9.0f}; + + // A present-but-empty bias supplies no value for any channel and must be rejected rather than + // indexing an empty vector. + OpTester test("ImageScaler"); + test.AddAttribute("scale", 2.0f); + test.AddAttribute("bias", std::vector{}); + test.AddInput("input", {N, C, H, W}, X); + test.AddOutput("output", {N, C, H, W}, std::vector(N * C * H * W, 0.0f)); + test.Run(OpTester::ExpectResult::kExpectFailure, + "Bias size (0) does not match the number of channels (2)", + {kTensorrtExecutionProvider}); +} + void MeanVarianceNormalizationAcrossChannels(bool across_channels, bool normalize_variance) { constexpr int64_t N = 2, C = 2, H = 2, W = 3; constexpr int64_t one = 1;