Use try-except for flash_attn

modeling_deepseek.py

@@ -48,7 +48,6 @@ from transformers.pytorch_utils import (
 from transformers.utils import (
     add_start_docstrings,
     add_start_docstrings_to_model_forward,
-    is_flash_attn_2_available,
     is_flash_attn_greater_or_equal_2_10,
     logging,
     replace_return_docstrings,
@@ -58,10 +57,11 @@ from .configuration_deepseek import DeepseekV2Config
 import torch.distributed as dist
 import numpy as np
 
-if is_flash_attn_2_available():
+try:
     from flash_attn import flash_attn_func, flash_attn_varlen_func
     from flash_attn.bert_padding import index_first_axis, pad_input, unpad_input  # noqa
-
+except ImportError:
+    pass
 
 # This makes `_prepare_4d_causal_attention_mask` a leaf function in the FX graph.
 # It means that the function will not be traced through and simply appear as a node in the graph.