diff options
| author |  avijit-nervana <avijit.chakraborty@intel.com> | 2018-09-05 21:58:53 -0700 |
|---|---|---|
| committer | avijit-nervana <avijit.chakraborty@intel.com> | 2018-09-05 21:58:53 -0700 |
| commit | df7fe4b73609a5ba2a5aa25a77882eb931cd0279 (patch) | |
| tree | 5620c738f3642c80564d806ee27097d16d31db7c /tensorflow/core/kernels/data/generator_dataset_op.cc | |
| parent | d972850a44ae624ad957b0dcc9d740e18f0cc10c (diff) | |
| parent | f4ae136265d3d3116a008b98ccf21d0791b878fd (diff) | |
Merge branch 'master' into avijit/add-cpu-backend
Diffstat (limited to 'tensorflow/core/kernels/data/generator_dataset_op.cc')
| -rw-r--r-- | tensorflow/core/kernels/data/generator_dataset_op.cc | 4 |
1 files changed, 4 insertions, 0 deletions
diff --git a/tensorflow/core/kernels/data/generator_dataset_op.cc b/tensorflow/core/kernels/data/generator_dataset_op.cc index ccee690d7e..ac5cc1b2c1 100644 --- a/tensorflow/core/kernels/data/generator_dataset_op.cc +++ b/tensorflow/core/kernels/data/generator_dataset_op.cc @@ -23,6 +23,7 @@ limitations under the License. #include "tensorflow/core/lib/random/random.h" namespace tensorflow { +namespace data { // See documentation in ../ops/dataset_ops.cc for a high-level // description of the following op. @@ -188,10 +189,13 @@ void GeneratorDatasetOp::MakeDataset(OpKernelContext* ctx, std::move(finalize_func), output_types_, output_shapes_); } +namespace { REGISTER_KERNEL_BUILDER(Name("GeneratorDataset").Device(DEVICE_CPU), GeneratorDatasetOp); REGISTER_KERNEL_BUILDER( Name("GeneratorDataset").Device(DEVICE_GPU).HostMemory("handle"), GeneratorDatasetOp); +} // namespace +} // namespace data } // namespace tensorflow |