try except import

modeling_llama_aqlm.py

@@ -56,8 +56,11 @@ from transformers.utils.import_utils import is_torch_fx_available
 from .configuration_llama_aqlm import LlamaConfig
 
 if is_flash_attn_2_available():
-    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
+    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:
+        pass
 
 
 # This makes `_prepare_4d_causal_attention_mask` a leaf function in the FX graph.