From 7ec1d78caf5ee6348f5388cad5dad2b16a1361da Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E9=BB=84=E5=AE=87=E6=89=AC?= <huang.yuyang@think-force.com>
Date: Sat, 13 Jul 2024 18:40:22 +0800
Subject: [PATCH] fix catdirectbatch

---
 src/devices/cuda/cudadevicebatch.cpp | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/devices/cuda/cudadevicebatch.cpp b/src/devices/cuda/cudadevicebatch.cpp
index 6e120dab..e9eb3393 100644
--- 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);
         }