From 4516282ccc621403ebf08a2975884220e3def1dc Mon Sep 17 00:00:00 2001
From: 12f23eddde <12f23eddde@gmail.com>
Date: Sat, 15 Feb 2025 22:04:45 +0800
Subject: [PATCH] Fix NoneType object has no attribute zero_

---
 ktransformers/models/custom_cache.py | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

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