From 32f62992a9fd7ddf180caf9a36d6d83da1a188d4 Mon Sep 17 00:00:00 2001
From: jiangfan06 <117341294+MuShangCC@users.noreply.github.com>
Date: Fri, 24 Mar 2023 10:37:57 +0800
Subject: [PATCH] [XPU] Fix argmax to adapt dynamic shape (#10126) (#10129)

---
 lite/kernels/xpu/argmax_compute.cc | 1 +
 1 file changed, 1 insertion(+)

diff --git a/lite/kernels/xpu/argmax_compute.cc b/lite/kernels/xpu/argmax_compute.cc
index a7b18f607aa..731cedc34b6 100644
--- 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>(