Fix onnx.If lowering with scalar condition tensor (#3846)

Fixes nod-ai/SHARK-ModelDev#696 (comment)
llvm · Nov 1, 2024 · 032a636 · 032a636
1 parent 25738b8
commit 032a636
Showing 1 changed file with 1 addition and 1 deletion.
diff --git 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");