diff --git a/lib/Conversion/TorchOnnxToTorch/DefaultDomainGtoP.cpp b/lib/Conversion/TorchOnnxToTorch/DefaultDomainGtoP.cpp
index a7f707cae..3de61f638 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");