[CustomDevice] fix bug in custom allocator (PaddlePaddle#56977)

BeingGod · Sep 9, 2023 · e36ef30 · e36ef30
1 parent 0ca6f53
commit e36ef30
Showing 1 changed file with 2 additions and 1 deletion.
diff --git 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;
 }