def _is_flex_attention_supported():
    """Check if flex attention is supported on the current device"""
    device = infer_device()
    return device in ["cuda"]
Check if flex attention is supported on the current device
test/transformers/test_flex_attention.py