Update modeling_llama.py

modeling_llama.py

@@ -1,7 +1,8 @@
 from typing import Optional, List, Union, Tuple
 
 import torch
-from transformers import LlamaConfig, LlamaModel, Cache, DynamicCache
+from torch import nn
+from transformers import LlamaModel, Cache, DynamicCache, LlamaForCausalLM
 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
@@ -143,9 +144,7 @@ class MightyLlamaModel(LlamaModel):
         )
 
 
-class MightyLlamaForCausalLM(LlamaForCausalLM):
-    config_class = LlamaConfig
-    
+class MightyLlamaForCausalLM(LlamaForCausalLM):    
     def __init__(self, config):
         super().__init__(config)
         self.model = MightyLlamaModel(config)