diff --git a/infer_helios.py b/infer_helios.py
index 23c730b..4c06c4d 100644
--- a/infer_helios.py
+++ b/infer_helios.py
@@ -214,7 +214,10 @@ def main():
         transformer = replace_rmsnorm_with_fp32(transformer)
         transformer = replace_all_norms_with_flash_norms(transformer)
         replace_rope_with_flash_rope()
-    transformer.set_attention_backend("_flash_3_hub")
+    try:
+        transformer.set_attention_backend("_flash_3_hub")
+    except Exception:
+        transformer.set_attention_backend("flash_hub")
 
     vae = AutoencoderKLWan.from_pretrained(
         args.base_model_path,