aboutsummaryrefslogtreecommitdiffhomepage
path: root/tensorflow/core/kernels/split_op.cc
diff options
context:
space:
mode:
authorGravatar Josh Levenberg <josh11b@tensorflow.org>2016-03-11 16:34:56 -0800
committerGravatar TensorFlower Gardener <gardener@tensorflow.org>2016-03-11 20:45:33 -0800
commitc5fc107300f7c89befcedf605622b571dd8b9f71 (patch)
treea7e897553ed885e57c94040d3865c1544c1937ae /tensorflow/core/kernels/split_op.cc
parent9917b53a41345cf30954a5a4b17985412f1acfde (diff)
Use the untemplated version of OpKernelContext::op_device_context() so
we don't need to #include tensorflow/core/common_runtime/gpu_device_context.h. Change: 117019921
Diffstat (limited to 'tensorflow/core/kernels/split_op.cc')
-rw-r--r--tensorflow/core/kernels/split_op.cc3
1 files changed, 1 insertions, 2 deletions
diff --git a/tensorflow/core/kernels/split_op.cc b/tensorflow/core/kernels/split_op.cc
index 8fae41e768..2912b0c094 100644
--- a/tensorflow/core/kernels/split_op.cc
+++ b/tensorflow/core/kernels/split_op.cc
@@ -27,7 +27,6 @@ limitations under the License.
#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
@@ -226,7 +225,7 @@ class SplitOpGPU : public SplitOpBase<GPUDevice, T> {
if (prefix_dim_size * split_dim_output_size * suffix_dim_size == 0) {
return;
}
- auto stream = context->op_device_context<GPUDeviceContext>()->stream();
+ auto stream = context->op_device_context()->stream();
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);