[Other] Fix compile error on linux aarch64 (PaddlePaddle#752)

Fix compile error on linux aarch64
YvonD · Nov 30, 2022 · 9d78b1d · 9d78b1d
1 parent e79dc86
commit 9d78b1d
Showing 1 changed file with 7 additions and 5 deletions.
diff --git a/fastdeploy/function/sort.cc b/fastdeploy/function/sort.cc
@@ -73,8 +73,9 @@ void SortKernel(const FDTensor& x, FDTensor* out, FDTensor* indices,
   axis = (axis < 0) ? (rank + axis) : axis;
   // Do full sort
   if (axis == -1 || axis + 1 == rank) {
-    const int64_t input_width = input_shape[rank - 1];
-    const int64_t input_height = x.Numel() / input_width;
+    int64_t numel = x.Numel();
+    int64_t input_width = input_shape[axis];
+    int64_t input_height = numel / input_width;
     FD_VISIT_INT_TYPES(indices_type, "FullSort", ([&] {
                          FullSort<T, data_t>(input_height, input_width, rank,
                                              &x, out, indices, descending);
@@ -93,8 +94,9 @@ void SortKernel(const FDTensor& x, FDTensor* out, FDTensor* indices,
 
     FDTensor trans_inp;
     Transpose(x, &trans_inp, trans);
-    const int64_t input_width = input_shape[axis];
-    const int64_t input_height = x.Numel() / input_width;
+    int64_t numel = x.Numel();
+    int64_t input_width = input_shape[axis];
+    int64_t input_height = numel / input_width;
     FD_VISIT_INT_TYPES(indices_type, "FullSort", ([&] {
                          FullSort<T, data_t>(input_height, input_width, rank,
                                              &trans_inp, out, indices,
@@ -115,4 +117,4 @@ void Sort(const FDTensor& x, FDTensor* out, FDTensor* indices, int axis,
 }
 
 }  // namespace function
-}  // namespace fastdeploy
+}  // namespace fastdeploy