diff --git a/onnxruntime/core/providers/cpu/tensor/mean_variance_normalization.h b/onnxruntime/core/providers/cpu/tensor/mean_variance_normalization.h index 7b41158ab5..924f380be6 100644 --- a/onnxruntime/core/providers/cpu/tensor/mean_variance_normalization.h +++ b/onnxruntime/core/providers/cpu/tensor/mean_variance_normalization.h @@ -105,10 +105,15 @@ class MeanVarianceNormalization_1 final : public MeanVarianceNormalization_0 if (!info.GetAttrs("axes", axes).IsOK()) { axes = {0, 2, 3}; } - if (find(axes.begin(), axes.end(), 1) != axes.end()) { + constexpr int64_t cross_channel_axes[] = {0, 1, 2, 3}; + constexpr int64_t batch_spatial_axes[] = {0, 2, 3}; + + if (std::equal(std::begin(axes), std::end(axes), std::begin(cross_channel_axes), std::end(cross_channel_axes))) { this->across_channels_ = true; - } else { + } else if (std::equal(std::begin(axes), std::end(axes), std::begin(batch_spatial_axes), std::end(batch_spatial_axes))) { this->across_channels_ = false; + } else { + ORT_THROW("MeanVarianceNormalization CPU EP only supports NHW and NCHW reduction for axes attribute."); } this->normalize_variance_ = 1; }