Skip to content

GitLab

Commit 4516282c authored by 12f23eddde's avatar 12f23eddde
Browse files

Fix NoneType object has no attribute zero_

parent 3c6035aa
Show whitespace changes
Inline Side-by-side
Showing with 2 additions and 1 deletion
ktransformers/models/custom_cache.py
View file @ 4516282c
...@@ -172,6 +172,7 @@ class StaticCache(transformers.StaticCache): ...@@ -172,6 +172,7 @@ class StaticCache(transformers.StaticCache):
for layer_idx in range(len(self.key_cache)): for layer_idx in range(len(self.key_cache)):
# In-place ops prevent breaking the static address # In-place ops prevent breaking the static address
self.key_cache[layer_idx].zero_() self.key_cache[layer_idx].zero_()
if self.value_cache[layer_idx] is not None:
self.value_cache[layer_idx].zero_() self.value_cache[layer_idx].zero_()
def get_max_cache_shape(self) -> Tuple[int, int, int, int]: def get_max_cache_shape(self) -> Tuple[int, int, int, int]:
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment