提交 2631077c 编写于 作者: L liaogang

Check Insufficient Alloc in tensor

上级 2200ff5e
......@@ -129,13 +129,16 @@ class Tensor {
virtual platform::Place place() const = 0;
};
template <typename T, typename PlaceType>
template <typename T, typename Place>
struct PlaceholderImpl : public Placeholder {
PlaceholderImpl(PlaceType place, size_t size)
: ptr_(static_cast<T*>(memory::Alloc(place, size)),
memory::PODDeleter<T, PlaceType>(place)),
place_(place),
size_(size) {}
PlaceholderImpl(Place place, size_t size)
: place_(place),
size_(size),
ptr_(static_cast<T*>(memory::Alloc(place, size)),
memory::PODDeleter<T, Place>(place)) {
PADDLE_ENFORCE(ptr_ != nullptr, "Insufficient %s memory to allocation.",
is_cpu_place(place_) ? "CPU" : "GPU");
}
virtual size_t size() const { return size_; }
virtual platform::Place place() const { return place_; }
......@@ -143,7 +146,7 @@ class Tensor {
virtual std::type_index type() const { return std::type_index(typeid(T)); }
/*! the pointer of memory block. */
std::unique_ptr<T, memory::PODDeleter<T, PlaceType>> ptr_;
std::unique_ptr<T, memory::PODDeleter<T, Place>> ptr_;
/*! the place of memory block. */
platform::Place place_;
......
Markdown is supported
0% .
You are about to add 0 people to the discussion. Proceed with caution.
先完成此消息的编辑!
想要评论请 注册