mirror of https://github.com/llvm/torch-mlir
Fix onnx.If lowering with scalar condition tensor (#3846)
Fixes https://github.com/nod-ai/SHARK-ModelDev/issues/696#issuecomment-2442016530pull/3847/head
parent
25738b8c19
commit
032a636c35
|
@ -180,7 +180,7 @@ void mlir::torch::onnx_c::populateDefaultDomainGtoP(
|
||||||
|
|
||||||
auto conditionType =
|
auto conditionType =
|
||||||
cast<Torch::ValueTensorType>(conditionTensor.getType());
|
cast<Torch::ValueTensorType>(conditionTensor.getType());
|
||||||
if (!conditionType || conditionType.getSizes().size() != 1)
|
if (!conditionType || conditionType.getSizes().size() > 1)
|
||||||
return rewriter.notifyMatchFailure(
|
return rewriter.notifyMatchFailure(
|
||||||
binder.op, "condition must have one single element per "
|
binder.op, "condition must have one single element per "
|
||||||
"https://onnx.ai/onnx/operators/onnx__If.html");
|
"https://onnx.ai/onnx/operators/onnx__If.html");
|
||||||
|
|
Loading…
Reference in New Issue