diff options
Diffstat (limited to 'tensorflow/core/kernels/data/prefetch_dataset_op.cc')
| -rw-r--r-- | tensorflow/core/kernels/data/prefetch_dataset_op.cc | 5 |
1 files changed, 5 insertions, 0 deletions
diff --git a/tensorflow/core/kernels/data/prefetch_dataset_op.cc b/tensorflow/core/kernels/data/prefetch_dataset_op.cc index 50efbcbe2a..a7a2935195 100644 --- a/tensorflow/core/kernels/data/prefetch_dataset_op.cc +++ b/tensorflow/core/kernels/data/prefetch_dataset_op.cc @@ -21,6 +21,7 @@ limitations under the License. #include "tensorflow/core/lib/core/error_codes.pb.h" namespace tensorflow { +namespace data { // See documentation in ../ops/dataset_ops.cc for a high-level // description of the following op. @@ -346,6 +347,7 @@ void PrefetchDatasetOp::MakeDataset(OpKernelContext* ctx, DatasetBase* input, *output = new Dataset(ctx, input, buffer_size); } +namespace { REGISTER_KERNEL_BUILDER(Name("PrefetchDataset").Device(DEVICE_CPU), PrefetchDatasetOp); REGISTER_KERNEL_BUILDER(Name("PrefetchDataset") @@ -354,4 +356,7 @@ REGISTER_KERNEL_BUILDER(Name("PrefetchDataset") .HostMemory("input_dataset") .HostMemory("handle"), PrefetchDatasetOp); +} // namespace + +} // namespace data } // namespace tensorflow |