Merge pull request #392 from TylunasLi/bug_fix_att_opt

修复非batch下CPU Attention算子取batch错误(#385)
ztxz16 · Dec 26, 2023 · 46a9918 · 46a9918
2 parents 4271f82 + c2fb785
commit 46a9918
Showing 1 changed file with 3 additions and 2 deletions.
diff --git a/src/devices/cpu/cpudevice.cpp b/src/devices/cpu/cpudevice.cpp
@@ -282,8 +282,9 @@ namespace fastllm {
         float *vd = (float*)v.cpuData;
         float *maskd = (datas.find("mask")->second && mask.dims.size() > 0) ? (float*)mask.cpuData : nullptr;
         float *od = (float*)output.cpuData;
-        int batch = intParams.find("q___batch")->second;
-        int maskStride = (datas.find("mask")->second) ? (mask.dims.size() == 3 ? mask.strides[0] : mask.Count(0)) : 0;
+        int batch = (maskd != nullptr && mask.dims.size() == 3) ? mask.dims[0] : 1; 
+        batch = intParams.find("mask___batch") != intParams.end() ? intParams.find("mask___batch")->second : batch;
+        int maskStride = (maskd != nullptr) ? (mask.dims.size() == 3 ? mask.strides[0] : mask.Count(0)) : 0;
         std::fill(od, od + output.Count(0), 0.0f);
         auto pool = GetPool();
         std::vector<std::future<void> > futures;