fix bug of trans_stack2Unsqueeze

senli123 · Jun 4, 2024 · 10ade66 · 10ade66
1 parent bacc8f9
commit 10ade66
Showing 1 changed file with 3 additions and 2 deletions.
diff --git a/tools/pnnx/src/pass_level6/trans_Stack2Unsqueeze.cpp b/tools/pnnx/src/pass_level6/trans_Stack2Unsqueeze.cpp
@@ -31,10 +31,10 @@ void trans_Stack2Unsqueeze(Graph& graph)
 
             if (op->type != "torch.stack")
                 continue;
-            matched = true;
             // get input num
             if( op->inputs.size() == 1)
             {
+                matched = true;
                 op->type = "torch.unsqueeze";
                 std::string str = op->name;
                 std::string from = "torch.stack";
@@ -47,8 +47,9 @@ void trans_Stack2Unsqueeze(Graph& graph)
                     str.replace(start_pos, from.length(), to);
                 }
                 op->name = str;
+                break;
             }
-            break;
+
         }
 
         if (!matched)