Update _modeling_kormo.py

_modeling_kormo.py

@@ -1,7 +1,6 @@
 from typing import Callable, List, Optional, Tuple, Union
 
 import torch
-import torch.utils.checkpoint ### ADD
 from torch import nn
 
 from transformers.activations import ACT2FN
@@ -17,13 +16,10 @@ from transformers.modeling_outputs import (
 )
 from transformers.modeling_rope_utils import ROPE_INIT_FUNCTIONS, dynamic_rope_update
 from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel
-# from transformers.pytorch_utils import ALL_LAYERNORM_LAYERS
 from transformers.processing_utils import Unpack
-from transformers.utils import auto_docstring, can_return_tuple, logging
+from transformers.utils import can_return_tuple, logging
 from ._configuration_kormo import KORMoConfig
 
-# from ._flash_attn3_doc import flash_attention_3_doc_forward
-# ALL_ATTENTION_FUNCTIONS._global_mapping.update({'flash_attention_3_doc': flash_attention_3_doc_forward})
 
 logger = logging.get_logger(__name__)
 
@@ -421,8 +417,6 @@ class KORMoModel(KORMoPreTrainedModel):
         )
 
     
-# class KwargsForCausalLM(FlashAttentionKwargs, LossKwargs): ...
-
 class KORMoForCausalLM(KORMoPreTrainedModel, GenerationMixin):
     _tied_weights_keys = ["lm_head.weight"]
     _tp_plan = {"lm_head": "colwise_rep"}