From 477516fcc5d5aa648475da7b77900dfcfc7bcb3e Mon Sep 17 00:00:00 2001
From: ronny1996 <ronny1996@163.com>
Date: Tue, 5 Sep 2023 09:29:30 +0000
Subject: [PATCH] [CustomDevice] fix bug in custom allocator

---
 paddle/fluid/memory/allocation/custom_allocator.cc | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/paddle/fluid/memory/allocation/custom_allocator.cc b/paddle/fluid/memory/allocation/custom_allocator.cc
index 37503105c8bc0..c8225a160428f 100644
--- a/paddle/fluid/memory/allocation/custom_allocator.cc
+++ b/paddle/fluid/memory/allocation/custom_allocator.cc
@@ -29,7 +29,8 @@ void CustomAllocator::FreeImpl(phi::Allocation* allocation) {
       platform::errors::PermissionDenied("CustomDevice memory is "
                                          "freed in incorrect device. "
                                          "This may be a bug"));
-
+  phi::DeviceManager::GetDeviceWithPlace(place_)->MemoryDeallocate(
+      allocation->ptr(), allocation->size());
   delete allocation;
 }