fix split-18

Signed-off-by: daquexian <[email protected]>
onnx · Nov 15, 2022 · 74fdf9c · 74fdf9c
1 parent 95b2161
commit 74fdf9c
Showing 1 changed file with 4 additions and 3 deletions.
diff --git a/onnxoptimizer/passes/eliminate_nop_split.h b/onnxoptimizer/passes/eliminate_nop_split.h
@@ -23,9 +23,10 @@ struct EliminateNopSplit final : public PredicateBasedPass {
   }
 
   bool patternMatchPredicate(Node* node) override {
-    return CheckKind(node, "Split") && node->input()->has_sizes() &&
-           node->hasAttribute(kaxis) && node->hasAttribute(ksplit) &&
-           node->outputs().size() == 1;
+    // TODO: support Split-18 where `split` is an input instead of an attribute
+    return CheckKind(node, "Split") && node->inputs().size() == 1 &&
+           node->input()->has_sizes() && node->hasAttribute(kaxis) &&
+           node->hasAttribute(ksplit) && node->outputs().size() == 1;
   }
 
   bool runTransform(Node* node, Graph& graph,