From 5c16b53e67482557b21f82c1a804352df3e76d85 Mon Sep 17 00:00:00 2001 From: zhink <771809832@qq.com> Date: Mon, 18 Dec 2023 20:53:20 +0800 Subject: [PATCH] fix cache kv len --- llm/predictor.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/llm/predictor.py b/llm/predictor.py index b4245e58b747..649285543b86 100644 --- a/llm/predictor.py +++ b/llm/predictor.py @@ -679,7 +679,7 @@ def __init__( model: PretrainedModel = None, tokenizer: PretrainedTokenizer = None, ): - self.cache_kvs_shape = model.get_cache_kvs_shape(model.config, config.batch_size) + self.cache_kvs_shape = model.get_cache_kvs_shape(model.config, config.batch_size, config.total_max_length) BasePredictor.__init__(self, config, tokenizer) InferencePredictorMixin.__init__(self, config, tokenizer) self.model = model