Update modeling_llama.py

modeling_llama.py

@@ -1,7 +1,7 @@
 from typing import Optional, List, Union, Tuple
 
 import torch
-from transformers import LlamaModel, Cache, DynamicCache
+from transformers import LlamaConfig, LlamaModel, Cache, DynamicCache
 from transformers.modeling_attn_mask_utils import _prepare_4d_causal_attention_mask_for_sdpa, \
     _prepare_4d_causal_attention_mask
 from transformers.modeling_outputs import BaseModelOutputWithPast
@@ -144,6 +144,8 @@ class MightyLlamaModel(LlamaModel):
 
 
 class MightyLlamaForCausalLM(LlamaForCausalLM):
+    config_class = LlamaConfig
+    
     def __init__(self, config):
         super().__init__(config)
         self.model = MightyLlamaModel(config)