fix: need_setup

UbiquitousLearning · Jul 21, 2024 · 00f29c5 · 00f29c5
1 parent da8909a
commit 00f29c5
Showing 1 changed file with 1 addition and 1 deletion.
diff --git 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] &