diff options
Diffstat (limited to 'tensorflow/core/kernels/data/writer_ops.cc')
| -rw-r--r-- | tensorflow/core/kernels/data/writer_ops.cc | 12 |
1 files changed, 6 insertions, 6 deletions
diff --git a/tensorflow/core/kernels/data/writer_ops.cc b/tensorflow/core/kernels/data/writer_ops.cc index 7bb2077b62..3f76695bb1 100644 --- a/tensorflow/core/kernels/data/writer_ops.cc +++ b/tensorflow/core/kernels/data/writer_ops.cc @@ -29,10 +29,10 @@ class ToTFRecordOp : public AsyncOpKernel { public: explicit ToTFRecordOp(OpKernelConstruction* ctx) : AsyncOpKernel(ctx), - background_worker_( - ctx->env(), - strings::StrCat("to_tf_record_op_", SanitizeThreadSuffix(name()))) { - } + thread_pool_(new thread::ThreadPool( + ctx->env(), ThreadOptions(), + strings::StrCat("to_tf_record__op_", SanitizeThreadSuffix(name())), + 1 /* num_threads */, false /* low_latency_hint */)) {} template <typename T> Status ParseScalarArgument(OpKernelContext* ctx, @@ -50,7 +50,7 @@ class ToTFRecordOp : public AsyncOpKernel { // The call to `iterator->GetNext()` may block and depend on an // inter-op thread pool thread, so we issue the call from the // owned thread pool. - background_worker_.Schedule([this, ctx, done]() { + thread_pool_->Schedule([this, ctx, done]() { string filename; OP_REQUIRES_OK_ASYNC( ctx, ParseScalarArgument<string>(ctx, "filename", &filename), done); @@ -97,7 +97,7 @@ class ToTFRecordOp : public AsyncOpKernel { } private: - BackgroundWorker background_worker_; + std::unique_ptr<thread::ThreadPool> thread_pool_; }; REGISTER_KERNEL_BUILDER(Name("DatasetToTFRecord").Device(DEVICE_CPU), |