fix

ztxz16 · Jul 15, 2024 · ad83823 · ad83823
1 parent 368f8c7
commit ad83823
Showing 1 changed file with 2 additions and 2 deletions.
diff --git 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.");
             }