From 3fe7a137b7a59e423a98bb57b2b199e63ba72320 Mon Sep 17 00:00:00 2001
From: "zhibin.xin" <zhibin.xin@sophgo.com>
Date: Fri, 12 Jul 2024 03:03:54 -0400
Subject: [PATCH] fix vit-b 4 batch prob

Change-Id: I5194338845ab3ec879cbce9521a77e855f675945
---
 lib/Dialect/Tpu/Transforms/LayerGroup/GroupMethod.cpp | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/lib/Dialect/Tpu/Transforms/LayerGroup/GroupMethod.cpp b/lib/Dialect/Tpu/Transforms/LayerGroup/GroupMethod.cpp
index 510361515..04164777d 100644
--- a/lib/Dialect/Tpu/Transforms/LayerGroup/GroupMethod.cpp
+++ b/lib/Dialect/Tpu/Transforms/LayerGroup/GroupMethod.cpp
@@ -105,7 +105,7 @@ static bool can_be_group_small_c(std::vector<Operation *> &group_ops) {
          shape[0] * shape[1] * shape[2] * shape[3] % Arch::NPU_NUM == 0)) {
       continue;
     }
-    if ((shape.size() == 3 && shape[0] > 1 && shape[1] == 197)) {
+    if ((shape.size() == 3 && shape[0] > 4 && shape[1] == 197)) {
       continue;
     }