Fix (#64092)

Author: co63oc

paddle/fluid/operators/math/beam_search.cc

@@ -69,15 +69,15 @@ class BeamSearchFunctor<phi::CPUContext, T> {
     // the output tensor shape should be [num_instances, 1]
     auto dims = common::make_ddim(
         std::vector<int64_t>({static_cast<int>(num_instances), 1}));
-    auto *selected_ids_data =
-        selected_ids->mutable_data<int64_t>(dims, platform::CPUPlace());
-    auto *selected_scores_data =
-        selected_scores->mutable_data<float>(dims, platform::CPUPlace());
+    selected_ids->Resize(dims);
+    auto *selected_ids_data = context.template Alloc<int64_t>(selected_ids);
+    selected_scores->Resize(dims);
+    auto *selected_scores_data = context.template Alloc<float>(selected_scores);
+    if (parent_idx != nullptr) {
+      parent_idx->Resize({static_cast<int64_t>(num_instances)});
+    }
     auto *parent_idx_data =
-        parent_idx
-            ? parent_idx->mutable_data<int>(
-                  {static_cast<int64_t>(num_instances)}, platform::CPUPlace())
-            : nullptr;
+        parent_idx ? context.template Alloc<int>(parent_idx) : nullptr;
 
     // fill in data
     std::vector<size_t> low_level;

paddle/fluid/operators/math/beam_search.cu

@@ -433,17 +433,18 @@ class BeamSearchFunctor<phi::GPUContext, T> {
     // Reserve a big enough memory.
     auto selected_dims =
         common::make_ddim({static_cast<int64_t>(num_seqs * beam_size), 1});
-    int64_t* selected_ids_data =
-        selected_ids->mutable_data<int64_t>(selected_dims, context.GetPlace());
+    selected_ids->Resize(selected_dims);
+    int64_t* selected_ids_data = context.template Alloc<int64_t>(selected_ids);
+    selected_scores->Resize(selected_dims);
     float* selected_scores_data =
-        selected_scores->mutable_data<float>(selected_dims, context.GetPlace());
+        context.template Alloc<float>(selected_scores);
+    if (parent_idx != nullptr) {
+      parent_idx->Resize({static_cast<int64_t>(num_seqs * beam_size)});
+    }
     int* parent_idx_data =
-        parent_idx ? parent_idx->mutable_data<int>(
-                         {static_cast<int64_t>(num_seqs * beam_size)},
-                         context.GetPlace())
-                   : nullptr;
+        parent_idx ? context.template Alloc<int>(parent_idx) : nullptr;
 
-    framework::LoD selected_lod(2);
+    phi::LoD selected_lod(2);
     selected_lod[0].assign(abs_lod[level].begin(), abs_lod[level].end());
     selected_lod[1].resize(scores->dims()[0] + 1);
     phi::MixVector<size_t> mix_vector(&selected_lod[1]);

paddle/fluid/operators/math/beam_search_xpu.cc

@@ -94,15 +94,15 @@ class BeamSearchFunctor<platform::XPUDeviceContext, T> {
     // the output tensor shape should be [num_instances, 1]
     auto dims = common::make_ddim(
         std::vector<int64_t>({static_cast<int>(num_instances), 1}));
-    auto *selected_ids_data =
-        selected_ids->mutable_data<int64_t>(dims, platform::CPUPlace());
-    auto *selected_scores_data =
-        selected_scores->mutable_data<float>(dims, platform::CPUPlace());
+    selected_ids->Resize(dims);
+    auto *selected_ids_data = context.template Alloc<int64_t>(selected_ids);
+    selected_scores->Resize(dims);
+    auto *selected_scores_data = context.template Alloc<float>(selected_scores);
+    if (parent_idx != nullptr) {
+      parent_idx->Resize({static_cast<int64_t>(num_instances)});
+    }
     auto *parent_idx_data =
-        parent_idx
-            ? parent_idx->mutable_data<int>(
-                  {static_cast<int64_t>(num_instances)}, platform::CPUPlace())
-            : nullptr;
+        parent_idx ? context.template Alloc<int>(parent_idx) : nullptr;
 
     // fill in data
     std::vector<size_t> low_level;

test/cpp/fluid/math/beam_search_test.cc

@@ -73,6 +73,9 @@ void TestBeamSearch() {
 
   auto* place = new Place();
   DeviceContext* context = new DeviceContext(*place);
+  context->SetAllocator(paddle::memory::allocation::AllocatorFacade::Instance()
+                            .GetAllocator(phi::CPUPlace())
+                            .get());
   if (paddle::platform::is_cpu_place(*place)) {
     PrepareCPUTensors(&ids, &scores, &pre_ids, &pre_scores);
   } else {