diff options
| author | Sam Gross <sgross@fb.com> | 2017-06-05 19:13:40 +0300 |
|---|---|---|
| committer | Soumith Chintala <soumith@gmail.com> | 2017-06-05 19:58:34 +0300 |
| commit | fc7d0950953a96a88d1eee7cb84b342361d38341 (patch) | |
| tree | 11de1ba4e7177c4b67397ab9d414632ef5a506fc | |
| parent | 9f5cefdc86b3e5bf0344cf236231c9a82b3fcda8 (diff) | |
Fix sharing of CUDA tensors on non-current devices
| -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 = { |