From 2980da559fb95fc6fc765eb86243c9f56654ffc8 Mon Sep 17 00:00:00 2001
From: duanjunwen <935724073@qq.com>
Date: Fri, 15 Nov 2024 10:26:30 +0000
Subject: [PATCH] [fix] fix test_lora

---
 colossalai/shardformer/policies/llama.py | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/colossalai/shardformer/policies/llama.py b/colossalai/shardformer/policies/llama.py
index d962057b1..b4a1f4bd8 100644
--- a/colossalai/shardformer/policies/llama.py
+++ b/colossalai/shardformer/policies/llama.py
@@ -506,8 +506,9 @@ class LlamaForSequenceClassificationPolicy(LlamaPolicy):
                 )
             }
             policy.update(new_item)
+        # TODO: test lora bug here
         # enable tp, replace layer to LinearWithGradAccum
-        else:
+        elif use_zbv:
             # add a new item for sequence classification
             new_item = {
                 LlamaForSequenceClassification: ModulePolicyDescription(