Store allocation ptr in vector (#65024)

eee4017 · lawrence910426 · web-flow · commit 5b2e91444539 · 2024-07-04T10:41:36.000+08:00
Co-authored-by: lawrence910426 &lt;lawu@nvidia.com&gt;
diff --git a/paddle/phi/kernels/funcs/segmented_array.h b/paddle/phi/kernels/funcs/segmented_array.h
@@ -112,7 +112,7 @@ struct ArraySetterBase {
                      void* src,
                      size_t num_bytes,
                      bool use_cuda_graph = false) {
-    allocation = phi::memory_utils::Alloc(
+    auto allocation = phi::memory_utils::Alloc(
         ctx.GetPlace(),
         num_bytes,
         phi::Stream(reinterpret_cast<phi::StreamId>(ctx.stream())));
@@ -129,10 +129,13 @@ struct ArraySetterBase {
                                        num_bytes,
                                        phi::gpuMemcpyHostToDevice,
                                        ctx.stream());
-    return allocation->ptr();
+
+    auto ptr = allocation->ptr();
+    allocations.emplace_back(std::move(allocation));
+    return ptr;
   }
 
-  phi::Allocator::AllocationPtr allocation{nullptr};
+  std::vector<phi::Allocator::AllocationPtr> allocations;
 };
 
 template <typename Context, typename T, SegmentedArraySize Size>
