fix catdirectbatch

ztxz16 · Jul 13, 2024 · 7ec1d78 · 7ec1d78
1 parent 112ce28
commit 7ec1d78
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/src/devices/cuda/cudadevicebatch.cpp b/src/devices/cuda/cudadevicebatch.cpp
@@ -308,7 +308,7 @@ namespace fastllm {
             dpitchs[b] = (input0Stride * unitSize);
             srcs[b] = (input1.cudaData);
             spitchs[b] = (input1Stride * unitSize);
-            widths[b] = (inner * unitSize);
+            widths[b] = (input1.dims[axis] * inner * unitSize);
             heights[b] = (outer);
         }