Skip to content

Commit b12c77d

Browse files
committed
Fix unittests
test=develop
1 parent 19e669a commit b12c77d

File tree

1 file changed

+9
-3
lines changed

1 file changed

+9
-3
lines changed

paddle/fluid/memory/allocation/allocator_facade.cc

Lines changed: 9 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,7 @@
1515
#include "paddle/fluid/memory/allocation/allocator.h"
1616
#include <gflags/gflags.h>
1717
#include <map>
18+
#include <string>
1819
#include <unordered_map>
1920
#include <vector>
2021
#include "paddle/fluid/memory/allocation/aligned_allocator.h"
@@ -209,6 +210,7 @@ class AllocatorFacadePrivate {
209210
for (int dev_id = 0; dev_id < platform::GetCUDADeviceCount(); ++dev_id) {
210211
places.emplace_back(platform::CUDAPlace(dev_id));
211212
}
213+
places.emplace_back(platform::CUDAPinnedPlace());
212214
#endif
213215
for (auto& p : places) {
214216
allocators_[p] = std::make_shared<LegacyAllocator>(p);
@@ -255,13 +257,17 @@ AllocatorFacade& AllocatorFacade::Instance() {
255257

256258
std::shared_ptr<Allocation> AllocatorFacade::AllocShared(
257259
const platform::Place& place, size_t size, Allocator::Attr attr) {
258-
return std::shared_ptr<Allocation>(
259-
m_->allocators_.at(place)->Allocate(size, attr).release(),
260-
AllocationDeleter());
260+
return std::shared_ptr<Allocation>(Alloc(place, size, attr).release(),
261+
AllocationDeleter());
261262
}
262263

263264
AllocationPtr AllocatorFacade::Alloc(const platform::Place& place, size_t size,
264265
Allocator::Attr attr) {
266+
auto it = m_->allocators_.find(place);
267+
if (it == m_->allocators_.end()) {
268+
throw BadAlloc(
269+
string::Sprintf("No such allocator for the place, %s", place));
270+
}
265271
return m_->allocators_.at(place)->Allocate(size, attr);
266272
}
267273

0 commit comments

Comments
 (0)