diff options
Diffstat (limited to 'tensorflow/core/kernels/data/prefetch_dataset_op.cc')
| -rw-r--r-- | tensorflow/core/kernels/data/prefetch_dataset_op.cc | 7 |
1 files changed, 6 insertions, 1 deletions
diff --git a/tensorflow/core/kernels/data/prefetch_dataset_op.cc b/tensorflow/core/kernels/data/prefetch_dataset_op.cc index 2bafb985ef..cc16108dce 100644 --- a/tensorflow/core/kernels/data/prefetch_dataset_op.cc +++ b/tensorflow/core/kernels/data/prefetch_dataset_op.cc @@ -357,7 +357,12 @@ class PrefetchDatasetOp : public UnaryDatasetOpKernel { REGISTER_KERNEL_BUILDER(Name("PrefetchDataset").Device(DEVICE_CPU), PrefetchDatasetOp); - +REGISTER_KERNEL_BUILDER(Name("PrefetchDataset") + .Device(DEVICE_GPU) + .HostMemory("buffer_size") + .HostMemory("input_dataset") + .HostMemory("handle"), + PrefetchDatasetOp); } // namespace } // namespace tensorflow |