Merge pull request Dao-AILab#60 from 201419/patch-1

fix typo in function mha_fwd
fabfish · Oct 17, 2022 · 1d0b41b · 1d0b41b
2 parents 52fb4b7 + ff07250
commit 1d0b41b
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/csrc/flash_attn/fmha_api.cpp b/csrc/flash_attn/fmha_api.cpp
@@ -216,7 +216,7 @@ mha_fwd(const at::Tensor &q,         // total_q x num_heads x head_size, total_q
     TORCH_CHECK(k.stride(-1) == 1);
     TORCH_CHECK(v.stride(-1) == 1);
     TORCH_CHECK(out.stride(-1) == 1);
-    TORCH_CHECK(cu_seqlens_k.is_contiguous());
+    TORCH_CHECK(cu_seqlens_q.is_contiguous());
     TORCH_CHECK(cu_seqlens_k.is_contiguous());
 
     const auto sizes = q.sizes();