fix

ztxz16 · Jul 17, 2024 · 3f594cf · 3f594cf
1 parent e18535c
commit 3f594cf
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/src/devices/cuda/cudadevicebatch.cpp b/src/devices/cuda/cudadevicebatch.cpp
@@ -255,7 +255,7 @@ 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[1]->dims.size();
+        int dimsLen = input1s[0]->dims.size();
         axis = (axis % dimsLen + dimsLen) % dimsLen;
         for (int i = 0; i < dimsLen && i < input0s[0]->dims.size(); i++) {
             if (i != axis) {