diff --git a/paddle/fluid/memory/allocation/allocator_facade.cc b/paddle/fluid/memory/allocation/allocator_facade.cc index b06ff1b485108fdc22ea8dc763fde18f7f955c3f..11c31df244e2a61caa2cd566a4b540bce6bdcb15 100644 --- a/paddle/fluid/memory/allocation/allocator_facade.cc +++ b/paddle/fluid/memory/allocation/allocator_facade.cc @@ -15,6 +15,7 @@ #include "paddle/fluid/memory/allocation/allocator.h" #include #include +#include #include #include #include "paddle/fluid/memory/allocation/aligned_allocator.h" @@ -209,6 +210,7 @@ class AllocatorFacadePrivate { for (int dev_id = 0; dev_id < platform::GetCUDADeviceCount(); ++dev_id) { places.emplace_back(platform::CUDAPlace(dev_id)); } + places.emplace_back(platform::CUDAPinnedPlace()); #endif for (auto& p : places) { allocators_[p] = std::make_shared(p); @@ -255,13 +257,17 @@ AllocatorFacade& AllocatorFacade::Instance() { std::shared_ptr AllocatorFacade::AllocShared( const platform::Place& place, size_t size, Allocator::Attr attr) { - return std::shared_ptr( - m_->allocators_.at(place)->Allocate(size, attr).release(), - AllocationDeleter()); + return std::shared_ptr(Alloc(place, size, attr).release(), + AllocationDeleter()); } AllocationPtr AllocatorFacade::Alloc(const platform::Place& place, size_t size, Allocator::Attr attr) { + auto it = m_->allocators_.find(place); + if (it == m_->allocators_.end()) { + throw BadAlloc( + string::Sprintf("No such allocator for the place, %s", place)); + } return m_->allocators_.at(place)->Allocate(size, attr); }