From 032a636c359456c80e5912eb53e1a2fe4d34f664 Mon Sep 17 00:00:00 2001
From: jinchen <49575973+jinchen62@users.noreply.github.com>
Date: Thu, 31 Oct 2024 20:34:50 -0700
Subject: [PATCH] Fix onnx.If lowering with scalar condition tensor (#3846)

Fixes
https://github.com/nod-ai/SHARK-ModelDev/issues/696#issuecomment-2442016530
---
 lib/Conversion/TorchOnnxToTorch/DefaultDomainGtoP.cpp | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/lib/Conversion/TorchOnnxToTorch/DefaultDomainGtoP.cpp b/lib/Conversion/TorchOnnxToTorch/DefaultDomainGtoP.cpp
index a7f707cae9bb..3de61f638fd7 100644
--- a/lib/Conversion/TorchOnnxToTorch/DefaultDomainGtoP.cpp
+++ b/lib/Conversion/TorchOnnxToTorch/DefaultDomainGtoP.cpp
@@ -180,7 +180,7 @@ void mlir::torch::onnx_c::populateDefaultDomainGtoP(
 
         auto conditionType =
             cast<Torch::ValueTensorType>(conditionTensor.getType());
-        if (!conditionType || conditionType.getSizes().size() != 1)
+        if (!conditionType || conditionType.getSizes().size() > 1)
           return rewriter.notifyMatchFailure(
               binder.op, "condition must have one single element per "
                          "https://onnx.ai/onnx/operators/onnx__If.html");