fix OOM error during resource allocation (#14444)

apache · Apr 5, 2019 · 09a4d64 · 09a4d64
1 parent c47d8c8
commit 09a4d64
Showing 1 changed file with 3 additions and 0 deletions.
diff --git a/src/resource.cc b/src/resource.cc
@@ -432,6 +432,9 @@ void Resource::get_cudnn_dropout_desc(
     // not initialized yet.
     size_t dropout_state_size;
     CUDNN_CALL(cudnnDropoutGetStatesSize(stream->dnn_handle_, &dropout_state_size));
+    // reserve GPU space
+    Storage::Get()->DirectFree(
+      Storage::Get()->Alloc(dropout_state_size, state_space->ctx));
     CUDNN_CALL(cudnnSetDropoutDescriptor(*dropout_desc, stream->dnn_handle_,
                                          dropout,
                                          state_space->GetSpace(dropout_state_size),