Commit 0babe27a authored by Davis King's avatar Davis King

Fixed race condition that could happen if set_size() was called while a cuda

kernel was still running.
parent 0f840db3
...@@ -81,6 +81,13 @@ namespace dlib ...@@ -81,6 +81,13 @@ namespace dlib
{ {
if (new_size == 0) if (new_size == 0)
{ {
if (device_in_use)
{
// Wait for any possible CUDA kernels that might be using our memory block to
// complete before we free the memory.
CHECK_CUDA(cudaStreamSynchronize(0));
device_in_use = false;
}
wait_for_transfer_to_finish(); wait_for_transfer_to_finish();
data_size = 0; data_size = 0;
host_current = true; host_current = true;
...@@ -91,6 +98,13 @@ namespace dlib ...@@ -91,6 +98,13 @@ namespace dlib
} }
else if (new_size != data_size) else if (new_size != data_size)
{ {
if (device_in_use)
{
// Wait for any possible CUDA kernels that might be using our memory block to
// complete before we free the memory.
CHECK_CUDA(cudaStreamSynchronize(0));
device_in_use = false;
}
wait_for_transfer_to_finish(); wait_for_transfer_to_finish();
data_size = new_size; data_size = new_size;
host_current = true; host_current = true;
......
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment