Skip to content

Commit

Permalink
Fix test
Browse files Browse the repository at this point in the history
  • Loading branch information
RAMitchell committed Oct 13, 2020
1 parent a97b042 commit 61c9a70
Show file tree
Hide file tree
Showing 2 changed files with 6 additions and 3 deletions.
5 changes: 3 additions & 2 deletions src/common/device_helpers.cuh
Expand Up @@ -412,10 +412,11 @@ struct XGBDefaultDeviceAllocatorImpl : XGBBaseDeviceAllocator<T> {
using other = XGBDefaultDeviceAllocatorImpl<U>; // NOLINT
};
pointer allocate(size_t n) { // NOLINT
pointer ptr = nullptr;
pointer ptr;
try {
ptr = SuperT::allocate(n);
} catch (const std::exception& e) {
dh::safe_cuda(cudaGetLastError());
} catch (const std::exception &e) {
ThrowOOMError(e.what(), n * sizeof(T));
}
GlobalMemoryLogger().RegisterAllocation(ptr.get(), n * sizeof(T));
Expand Down
4 changes: 3 additions & 1 deletion tests/cpp/common/test_device_helpers.cu
Expand Up @@ -161,6 +161,8 @@ TEST(Allocator, OOM) {
auto size = dh::AvailableMemory(0) * 4;
ASSERT_THROW({dh::caching_device_vector<char> vec(size);}, dmlc::Error);
ASSERT_THROW({dh::device_vector<char> vec(size);}, dmlc::Error);
// Clear last error so we don't fail subsequent tests
cudaGetLastError();
}
} // namespace common
} // namespace xgboost
} // namespace xgboost

0 comments on commit 61c9a70

Please sign in to comment.