#if GOOGLE_CUDA #define EIGEN_USE_GPU #include #include "tensorflow/core/kernels/split_op.h" #include "tensorflow/core/framework/register_types.h" #include "tensorflow/core/framework/tensor_types.h" namespace tensorflow { namespace functor { template void Split::operator()( const Device& d, typename TTypes::Tensor output, typename TTypes::ConstTensor input, const Eigen::DSizes& slice_indices, const Eigen::DSizes& slice_sizes) { output.device(d) = input.slice(slice_indices, slice_sizes); } #define DEFINE_GPU_KERNELS(T) template struct Split; TF_CALL_GPU_NUMBER_TYPES(DEFINE_GPU_KERNELS); } // namespace functor } // namespace tensorflow #endif // GOOGLE_CUDA