fix get N bug

57b856b3 · ltqin · 8afad0f6 · 57b856b3
Commit 57b856b3 authored Dec 30, 2022 by ltqin
Hide whitespace changes
Inline Side-by-side

Showing with 1 addition and 1 deletion

example/32_batched_gemm_scale_softmax_gemm/batched_multihead_attention_backward_fp16.cpp ...oftmax_gemm/batched_multihead_attention_backward_fp16.cpp +1 -1

No files found.
--- a/example/32_batched_gemm_scale_softmax_gemm/batched_multihead_attention_backward_fp16.cpp
+++ b/example/32_batched_gemm_scale_softmax_gemm/batched_multihead_attention_backward_fp16.cpp
@@ -210,7 +210,7 @@ void run_attention_fwd_host(const TensorQ& q_g_m_k,

    // masking
 #if USING_MASK
-    auto N          = s_g_m_n.GetLengths()[1];
+    auto N          = s_g_m_n.GetLengths()[2];
    const auto mask = DeviceGemmInstance::C0MatrixMask(N);
    s_g_m_n.ForEach([&](auto& self, auto idx) {
        if(mask.IsMaskedElement(idx[1], idx[2]))