From 00f29c5646cd265a776d9bfe2ff75874d2f7e0cb Mon Sep 17 00:00:00 2001
From: yirongjie <yirj0809@gmail.com>
Date: Sun, 21 Jul 2024 06:43:25 +0000
Subject: [PATCH] fix: need_setup

---
 src/Module.hpp | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/Module.hpp b/src/Module.hpp
index c0470e01..a696220a 100644
--- a/src/Module.hpp
+++ b/src/Module.hpp
@@ -145,7 +145,7 @@ class Module {
                 if(input.batch() == 0){
                     Tensor::gph_[input.name()] = input;
                 }
-                if(input.sequence()!=1 && !last_shape_bshd_.empty()){
+                if(inputs[0].sequence()!=1 && !last_shape_bshd_.empty()){
                     // if LLM/VLLM model, the `need_setup` should be `true`
                     if(input.batch() == last_shape_bshd_[i][0] & 
                         input.sequence() == last_shape_bshd_[i][1] &