Merge branch 'develop' of https://github.com/Vvsmile/Paddle into develop

PaddlePaddle · Jun 20, 2023 · 7bf3ed9 · 7bf3ed9
2 parents 7f770c8 + 49a45f7
commit 7bf3ed9
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/python/paddle/nn/functional/flash_attention.py b/python/paddle/nn/functional/flash_attention.py
@@ -81,7 +81,7 @@ def _math_attention(
 
 
 def _select_sdp_cuda(head_dim):
-    if head_dim < 128:
+    if head_dim <= 128:
         return "flash_attn"
     else:
         return "mem_efficient"