fixed cache over-alloc bug (#17)

- fixed cache over-alloc bug (8ca4513e59cb96cc89023824b75a18dda653644d)


Co-authored-by: Tomer Ronen <[email protected]>

variable_cache.py

@@ -83,6 +83,7 @@ class VariableCache(Cache_4_44_2, Cache):
         if attention_config.no_op or attention_config.replace_with_linear:
             return None
         config = deepcopy(self.config)
+        config.num_hidden_layers = 1
         config.num_key_value_heads = self.config.num_attention_heads // attention_config.n_heads_in_group
         return StaticCache(config, self.max_batch_size, self.max_cache_len, device, self.dtype)