fix fc fuse proble (#36568)

PaddlePaddle · Oct 20, 2021 · fc5db55 · fc5db55
1 parent 3f2d6a3
commit fc5db55
Showing 1 changed file with 6 additions and 1 deletion.
diff --git a/paddle/fluid/framework/ir/fc_fuse_pass.cc b/paddle/fluid/framework/ir/fc_fuse_pass.cc
@@ -51,7 +51,12 @@ FCFusePass::FCFusePass() {
       .IsTensor()
       .End()
       .AddAttr("axis")
-      .IsNumGE(1)
+      .IsNumMatch<int>([](int axis) -> bool {
+        if (axis == -1 || axis >= 1) {
+          return true;
+        }
+        return false;
+      })
       .End();
 
   AddOpCompat(OpCompat("relu"))