[XPU] Fix argmax to adapt dynamic shape (#10126) (#10129)

PaddlePaddle · Mar 24, 2023 · 32f6299 · 32f6299
1 parent 2072742
commit 32f6299
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/lite/kernels/xpu/argmax_compute.cc b/lite/kernels/xpu/argmax_compute.cc
@@ -54,6 +54,7 @@ void ArgmaxCompute::Run() {
     CHECK_EQ(r, 0);
   } else if (param.dtype == 2) {
     // int32
+    out_int64_xpu_guard_->Reserve(out->numel() * sizeof(int64_t));
     int64_t* out_int64_data =
         reinterpret_cast<int64_t*>(out_int64_xpu_guard_->addr_);
     int r = xdnn::argmax<float, int64_t>(