From 7e7e2b7bf8aa4631313a1701dd2826ccb955de64 Mon Sep 17 00:00:00 2001 From: ronnywang Date: Wed, 23 Aug 2023 10:07:05 +0800 Subject: [PATCH] [CustomDevice] fix default stream allocator error (#56536) --- .../allocation/stream_safe_custom_device_allocator.cc | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/paddle/fluid/memory/allocation/stream_safe_custom_device_allocator.cc b/paddle/fluid/memory/allocation/stream_safe_custom_device_allocator.cc index 0658dacc986..d0faf544e00 100644 --- 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 #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::DeviceContextPool::Instance().Get(place_)) + ->stream()); } stream_safe_cuda_allocation->MarkAsWillBeFreed(); if (stream_safe_cuda_allocation->CanBeFreed()) { -- GitLab