From e603e78627ac2765301e0f8e9a5f76f8fb2fe9ec Mon Sep 17 00:00:00 2001
From: Dmitri Smirnov <yuslepukhin@users.noreply.github.com>
Date: Wed, 6 Dec 2023 21:04:18 -0800
Subject: [PATCH] Enforce If condition size == 1 (#18733)

### Description
<!-- Describe your changes. -->

### Motivation and Context
https://github.com/microsoft/onnxruntime/issues/18549
---
 onnxruntime/core/providers/cpu/controlflow/if.cc | 7 ++++++-
 1 file changed, 6 insertions(+), 1 deletion(-)

diff --git a/onnxruntime/core/providers/cpu/controlflow/if.cc b/onnxruntime/core/providers/cpu/controlflow/if.cc
index a5fe3f02b2..51d2fc8291 100644
--- a/onnxruntime/core/providers/cpu/controlflow/if.cc
+++ b/onnxruntime/core/providers/cpu/controlflow/if.cc
@@ -248,7 +248,12 @@ Status If::Compute(OpKernelContext* ctx) const {
 
   auto ctx_internal = static_cast<OpKernelContextInternal*>(ctx);
 
-  auto condition = *ctx->Input<Tensor>(0)->Data<bool>();
+  const auto& condition_tensor = *ctx->Input<Tensor>(0);
+
+  ORT_RETURN_IF_NOT(condition_tensor.Shape().Size() == 1,
+                    "If nodes condition input must have exactly one element");
+
+  auto condition = *condition_tensor.Data<bool>();
 
   auto attribute = condition ? "then_branch" : "else_branch";
   auto* session_state = ctx_internal->SubgraphSessionState(attribute);