fix(kernel):修复transpose info存在的bug

InfiniTensor · Jan 31, 2024 · 7db54c1 · 7db54c1
1 parent 3f5b0ee
commit 7db54c1
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/src/04kernel/src/attributes/transpose_info.cc b/src/04kernel/src/attributes/transpose_info.cc
@@ -35,7 +35,7 @@ namespace refactor::kernel {
                 }
             }
         }
-        if (rank == 0) {
+        if (rank <= 1) {
             dims = {{1, 1}};
             blockSize *= blockCount;
             blockCount = 1;
@@ -81,6 +81,7 @@ namespace refactor::kernel {
             perm.pop_back();
             --rank;
         }
+
         // 计算 stride
         struct StrideI {
             dim_t strideI;