diff --git a/ktransformers/models/custom_cache.py b/ktransformers/models/custom_cache.py index 8a11f1f..45aa436 100644 --- a/ktransformers/models/custom_cache.py +++ b/ktransformers/models/custom_cache.py @@ -172,7 +172,8 @@ class StaticCache(transformers.StaticCache): for layer_idx in range(len(self.key_cache)): # In-place ops prevent breaking the static address self.key_cache[layer_idx].zero_() - self.value_cache[layer_idx].zero_() + if self.value_cache[layer_idx] is not None: + self.value_cache[layer_idx].zero_() def get_max_cache_shape(self) -> Tuple[int, int, int, int]: """Returns the maximum shape of the cache."""