fix torch_attn import (#355)

xdit-project · Nov 20, 2024 · 20b1a60 · 20b1a60
1 parent 9b61d76
commit 20b1a60
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/xfuser/core/long_ctx_attention/ulysses/attn_layer.py b/xfuser/core/long_ctx_attention/ulysses/attn_layer.py
@@ -10,11 +10,11 @@
 from yunchang.comm.all_to_all import SeqAllToAll4D
 try:
     # yunchang > 0.4.0
-    from yunchang.ulysses.attn_layer import torch_attn
+    from yunchang.kernels.attention import torch_attn
 except:
     print(f"detect you are not use the latest yunchang. Please install yunchang>=0.4.0")
     try:
-        from yunchang.kernels.attention import torch_attn
+        from yunchang.ulysses.attn_layer import torch_attn
     except:
         raise ImportError(f"yunchang import torch_attn error")