RYDE-WORK/ktransformers
1
0
Fork 0
mirror of https://github.com/RYDE-WORK/ktransformers.git synced 2026-01-19 12:43:16 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

Merge pull request #330 from hrz6976/fix-nonetype

Browse Source 
thanks..., I was about to submit and found that you had already modified it. Thank you for your contribution
This commit is contained in:
wang jiahao 2025-02-15 22:50:51 +08:00 committed by GitHub
commit ae8da019c1
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 2 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

3
ktransformers/models/custom_cache.py
View File

@ -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."""