Update vision.py

vision.py

@@ -34,10 +34,16 @@ from .configuration_vmistral import VMistralVisionConfig
 
 logger = logging.get_logger(__name__)
 
-
+# is_flash_attn_2_available() checks if flash_attn is available, if yes, it imports it
 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
+# if flash_attn is not available, it raises an ImportError, intimating the user to install flash_attn
+else:
+    raise ImportError(
+        "Flash Attention 2.0 is not available. Please install flash-attn>=2.1.0: "
+        "`pip install flash-attn>=2.1.0` or use an alternative attention implementation."
+    )
 
 
 # Copied from transformers.models.llama.modeling_llama._get_unpad_data