update flash attn select (#54630) (#54716)

PaddlePaddle · Jun 19, 2023 · feff99f · feff99f
1 parent 570daa1
commit feff99f
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"