diff options
| -rw-r--r-- | lib/THC/THCAllocator.c | 15 |
1 files changed, 14 insertions, 1 deletions
diff --git a/lib/THC/THCAllocator.c b/lib/THC/THCAllocator.c index 9ff447d..830363a 100644 --- a/lib/THC/THCAllocator.c +++ b/lib/THC/THCAllocator.c @@ -32,7 +32,20 @@ static cudaError_t THCIpcAllocator_malloc(void* ctx, void** devPtr, size_t size, static cudaError_t THCIpcAllocator_free(void* ctx, void* devPtr) { - return cudaIpcCloseMemHandle(devPtr); + cudaError_t err; + int prev_device; + int device = (int)(long)ctx; + + err = cudaGetDevice(&prev_device); + if (err != cudaSuccess) { return err; } + + err = cudaSetDevice(device); + if (err != cudaSuccess) { return err; } + + err = cudaIpcCloseMemHandle(devPtr); + + cudaSetDevice(prev_device); + return err; } THCDeviceAllocator THCIpcAllocator = { |