| 102 | } |
| 103 | |
| 104 | void SyncedMemory::set_gpu_data(void* data) { |
| 105 | #ifndef CPU_ONLY |
| 106 | CHECK(data); |
| 107 | if (own_gpu_data_) { |
| 108 | int initial_device; |
| 109 | cudaGetDevice(&initial_device); |
| 110 | if (gpu_device_ != -1) { |
| 111 | CUDA_CHECK(cudaSetDevice(gpu_device_)); |
| 112 | } |
| 113 | CUDA_CHECK(cudaFree(gpu_ptr_)); |
| 114 | cudaSetDevice(initial_device); |
| 115 | } |
| 116 | gpu_ptr_ = data; |
| 117 | head_ = HEAD_AT_GPU; |
| 118 | own_gpu_data_ = false; |
| 119 | #else |
| 120 | NO_GPU; |
| 121 | #endif |
| 122 | } |
| 123 | |
| 124 | void* SyncedMemory::mutable_cpu_data() { |
| 125 | to_cpu(); |