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
{
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();
data_size = 0;
host_current = true;
......@@ -91,6 +98,13 @@ namespace dlib
}
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();
data_size = new_size;
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