From 54cdb6af71e50c002d32b5ac8324c47c185dfc5c Mon Sep 17 00:00:00 2001 From: Scott McKay Date: Thu, 4 Mar 2021 20:42:40 +1000 Subject: [PATCH] Add check that the first 2 Loop subgraph inputs have an shape (could be explicit or inferred) as we need to know the rank the subgraph expects. Other inputs to the subgraph are more opaque so we can just pass them through. (#6891) --- onnxruntime/core/providers/cpu/controlflow/loop.cc | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/onnxruntime/core/providers/cpu/controlflow/loop.cc b/onnxruntime/core/providers/cpu/controlflow/loop.cc index 23d8ab6427..a2b7c4b961 100644 --- a/onnxruntime/core/providers/cpu/controlflow/loop.cc +++ b/onnxruntime/core/providers/cpu/controlflow/loop.cc @@ -399,6 +399,12 @@ Status LoopImpl::Initialize() { auto& subgraph_inputs = info_.subgraph.GetInputs(); + // we need to know if the subgraph expects a rank 0 or rank 1 value for these, so a shape is required. + ORT_RETURN_IF(subgraph_inputs[0]->Shape() == nullptr, "Loop subgraph input 0 has unknown shape: ", + subgraph_inputs[0]->Name()); + ORT_RETURN_IF(subgraph_inputs[1]->Shape() == nullptr, "Loop subgraph input 1 has unknown shape: ", + subgraph_inputs[1]->Name()); + auto iter_num_rank = subgraph_inputs[0]->Shape()->dim_size(); auto condition_rank = subgraph_inputs[1]->Shape()->dim_size();