Fix test

dmlc · Oct 13, 2020 · a9d6135 · a9d6135
1 parent a97b042
commit a9d6135
Showing 1 changed file with 3 additions and 2 deletions.
diff --git a/src/common/device_helpers.cuh b/src/common/device_helpers.cuh
@@ -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));