Skip to content

Commit

Permalink
Enforce If condition size == 1 (#18733)
Browse files Browse the repository at this point in the history
### Description
<!-- Describe your changes. -->

### Motivation and Context
#18549
  • Loading branch information
yuslepukhin authored Dec 7, 2023
1 parent 9479ba5 commit e603e78
Showing 1 changed file with 6 additions and 1 deletion.
7 changes: 6 additions & 1 deletion onnxruntime/core/providers/cpu/controlflow/if.cc
Original file line number Diff line number Diff line change
Expand Up @@ -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);
Expand Down

0 comments on commit e603e78

Please sign in to comment.