Update modeling_llama.py

modeling_llama.py

@@ -294,7 +294,7 @@ class LlamaAttention(nn.Module):
             key_states = torch.cat([past_key_value[0], key_states], dim=2)
             value_states = torch.cat([past_key_value[1], value_states], dim=2)
 
-        past_key_value = (cache_key_states, value_states) if use_cache else None
+        past_key_value = (key_states, value_states) if use_cache else None
 
         use_flashattn =  self.config.use_flashattn and is_flash_attn_available()