[CustomDevice] fix default stream allocator error (#56536)

ronny1996 · web-flow · commit 7e7e2b7bf8aa · 2023-08-23T10:07:05.000+08:00
diff --git a/paddle/fluid/memory/allocation/stream_safe_custom_device_allocator.cc b/paddle/fluid/memory/allocation/stream_safe_custom_device_allocator.cc
@@ -16,6 +16,7 @@
 #include <thread>
 
 #include "paddle/fluid/platform/profiler/event_tracing.h"
+#include "paddle/phi/backends/context_pool.h"
 
 namespace paddle {
 namespace memory {
@@ -179,7 +180,11 @@ void StreamSafeCustomDeviceAllocator::FreeImpl(phi::Allocation* allocation) {
 
   VLOG(8) << "Try free allocation " << stream_safe_cuda_allocation->ptr();
   if (!stream_safe_cuda_allocation->GetOwningStream()) {
-    stream_safe_cuda_allocation->SetOwningStream(default_stream_);
+    stream_safe_cuda_allocation->SetOwningStream(
+        default_stream_ ? default_stream_
+                        : reinterpret_cast<phi::CustomContext*>(
+                              phi::DeviceContextPool::Instance().Get(place_))
+                              ->stream());
   }
   stream_safe_cuda_allocation->MarkAsWillBeFreed();
   if (stream_safe_cuda_allocation->CanBeFreed()) {
