From ad83823bdc9b14e954e1d59504b4373030f99b2f 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: Mon, 15 Jul 2024 10:13:06 +0800
Subject: [PATCH] fix

---
 src/devices/cuda/cudadevicebatch.cpp | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/src/devices/cuda/cudadevicebatch.cpp b/src/devices/cuda/cudadevicebatch.cpp
index e9eb3393..db381f75 100644
--- a/src/devices/cuda/cudadevicebatch.cpp
+++ b/src/devices/cuda/cudadevicebatch.cpp
@@ -255,9 +255,9 @@ namespace fastllm {
                             "CatDirect error: inputs should use same device.\n");
         AssertInFastLLM(input0s[0]->dims.size() == 0 || input0s[0]->dims.size() == input1s[0]->dims.size(),
                             "Cat Error: input's shape's size should be same.\n");
-        int dimsLen = input0s[0]->dims.size();
+        int dimsLen = input0s[1]->dims.size();
         axis = (axis % dimsLen + dimsLen) % dimsLen;
-        for (int i = 0; i < dimsLen; i++) {
+        for (int i = 0; i < dimsLen && i < input0s[0]->dims.size(); i++) {
             if (i != axis) {
                 AssertInFastLLM(input0s[0]->dims[i] == input1s[0]->dims[i], "Cat Error: input's shape doesn't match.");
             }