refactor(evo1): use transformers flash attention probe (#4013)

Co-authored-by: Martino Russi <77496684+nepyope@users.noreply.github.com>

src/lerobot/policies/evo1/internvl3_embedder.py

@@ -27,9 +27,11 @@ from lerobot.utils.import_utils import _transformers_available, require_package
 
 if TYPE_CHECKING or _transformers_available:
     from transformers import AutoModel, AutoTokenizer
+    from transformers.utils import is_flash_attn_2_available
 else:
     AutoModel = None
     AutoTokenizer = None
+    is_flash_attn_2_available = None
 
 IMAGENET_MEAN = (0.485, 0.456, 0.406)
 IMAGENET_STD = (0.229, 0.224, 0.225)
@@ -135,9 +137,13 @@ class InternVL3Embedder(nn.Module):
                 raise ValueError(f"Unsupported EVO1 vlm_dtype '{model_dtype}'") from exc
         self.model_dtype = model_dtype
 
-        attn_implementation = "flash_attention_2" if (use_flash_attn and _flash_attn_available()) else "eager"
+        attn_implementation = (
+            "flash_attention_2" if (use_flash_attn and is_flash_attn_2_available()) else "eager"
+        )
         if use_flash_attn and attn_implementation == "eager":
-            logger.warning("flash_attn is not installed. Falling back to eager attention.")
+            logger.warning(
+                "Flash Attention 2 is unavailable on this runtime. Falling back to eager attention."
+            )
 
         self.model = AutoModel.from_pretrained(
             model_name,
@@ -359,11 +365,3 @@ class InternVL3Embedder(nn.Module):
     @property
     def device(self) -> torch.device:
         return next(self.model.parameters()).device
-
-
-def _flash_attn_available() -> bool:
-    try:
-        import flash_attn  # noqa: F401
-    except ModuleNotFoundError:
-        return False
-    return True