fix: bug

duterscmy · Dec 12, 2024 · f371740 · f371740
1 parent 84dbe01
commit f371740
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/cd-moe/modeling_deepseek.py b/cd-moe/modeling_deepseek.py
@@ -684,7 +684,7 @@ def forward(
         if past_key_value is not None:
             # print(self.layer_idx)
             cache_kwargs = {"sin": sin, "cos": cos}  # Specific to RoPE models
-            cache_idx = self.layer_idx + 1 if self.layer_idx > 0 else 0
+            cache_idx = self.layer_idx if self.layer_idx > 0 else 0
             key_states, value_states = past_key_value.update(
                 key_states, value_states, cache_idx, cache_kwargs)