diff options
| author |  Vijay Vasudevan <vrv@google.com> | 2016-02-25 11:10:51 -0800 |
|---|---|---|
| committer |  TensorFlower Gardener <gardener@tensorflow.org> | 2016-02-25 11:16:20 -0800 |
| commit | 13d7f5203454d328c6cc6cb048cc3682cc53fe9e (patch) | |
| tree | 7afed07e0f7edc1c945d6e8d7498dde9328bef76 | |
| parent | 86e93febaa2f349f2fbf733d6895a69e87ec578e (diff) | |
TensorFlow: make split_op not use internal header library for callback,
since this breaks the build on GPU.
Change: 115582331
| -rw-r--r-- | tensorflow/core/kernels/split_op.cc | 10 |
1 files changed, 4 insertions, 6 deletions
diff --git a/tensorflow/core/kernels/split_op.cc b/tensorflow/core/kernels/split_op.cc index 8fae41e768..5dc6cbb015 100644 --- a/tensorflow/core/kernels/split_op.cc +++ b/tensorflow/core/kernels/split_op.cc @@ -26,7 +26,6 @@ limitations under the License. #include "tensorflow/core/lib/core/status.h" #include "tensorflow/core/lib/gtl/array_slice.h" #if GOOGLE_CUDA -#include "tensorflow/core/common_runtime/gpu/gpu_event_mgr.h" #include "tensorflow/core/common_runtime/gpu_device_context.h" #include "tensorflow/core/platform/stream_executor.h" #endif // GOOGLE_CUDA @@ -230,11 +229,10 @@ class SplitOpGPU : public SplitOpBase<GPUDevice, T> { perftools::gputools::DeviceMemoryBase output_ptrs_base{ output_ptrs_on_gpu.flat<int8>().data(), static_cast<uint64>(num_split)}; TensorReference tensor_ref(output_ptrs_on_host); - stream->ThenMemcpy(&output_ptrs_base, - output_ptrs_on_host.flat<int8>().data(), - output_ptrs_total_bytes); - context->device()->tensorflow_gpu_device_info()->event_mgr->ThenExecute( - stream, [tensor_ref]() { tensor_ref.Unref(); }); + stream + ->ThenMemcpy(&output_ptrs_base, output_ptrs_on_host.flat<int8>().data(), + output_ptrs_total_bytes) + .ThenDoHostCallback([tensor_ref]() { tensor_ref.Unref(); }); SplitOpGPULaunch<T>().Run( context->eigen_device<GPUDevice>(), input.flat<T>().data(), num_split, prefix_dim_size, split_dim_size, suffix_dim_size, |