From b12c77dae258480db23b4d98c44e61026a630330 Mon Sep 17 00:00:00 2001
From: Yu Yang <yuyang18@baidu.com>
Date: Mon, 19 Nov 2018 09:35:07 +0800
Subject: [PATCH] Fix unittests

test=develop
---
 paddle/fluid/memory/allocation/allocator_facade.cc | 12 +++++++++---
 1 file changed, 9 insertions(+), 3 deletions(-)

diff --git a/paddle/fluid/memory/allocation/allocator_facade.cc b/paddle/fluid/memory/allocation/allocator_facade.cc
index b06ff1b4851..11c31df244e 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 <gflags/gflags.h>
 #include <map>
+#include <string>
 #include <unordered_map>
 #include <vector>
 #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<LegacyAllocator>(p);
@@ -255,13 +257,17 @@ AllocatorFacade& AllocatorFacade::Instance() {
 
 std::shared_ptr<Allocation> AllocatorFacade::AllocShared(
     const platform::Place& place, size_t size, Allocator::Attr attr) {
-  return std::shared_ptr<Allocation>(
-      m_->allocators_.at(place)->Allocate(size, attr).release(),
-      AllocationDeleter());
+  return std::shared_ptr<Allocation>(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);
 }
 
-- 
GitLab