diff options
Diffstat (limited to 'tensorflow/core/kernels/data/writer_ops.cc')
-rw-r--r-- | tensorflow/core/kernels/data/writer_ops.cc | 11 |
1 files changed, 6 insertions, 5 deletions
diff --git a/tensorflow/core/kernels/data/writer_ops.cc b/tensorflow/core/kernels/data/writer_ops.cc index 80d9a5b867..1c49874a6a 100644 --- a/tensorflow/core/kernels/data/writer_ops.cc +++ b/tensorflow/core/kernels/data/writer_ops.cc @@ -70,20 +70,21 @@ class ToTFRecordOp : public AsyncOpKernel { DatasetBase* dataset; OP_REQUIRES_OK_ASYNC( ctx, GetDatasetFromVariantTensor(ctx->input(0), &dataset), done); - IteratorContext iter_ctx = dataset::MakeIteratorContext(ctx); std::unique_ptr<IteratorBase> iterator; OP_REQUIRES_OK_ASYNC( ctx, - dataset->MakeIterator(&iter_ctx, "ToTFRecordOpIterator", &iterator), + dataset->MakeIterator(IteratorContext(ctx), "ToTFRecordOpIterator", + &iterator), done); std::vector<Tensor> components; components.reserve(dataset->output_dtypes().size()); bool end_of_sequence; do { - OP_REQUIRES_OK_ASYNC( - ctx, iterator->GetNext(&iter_ctx, &components, &end_of_sequence), - done); + OP_REQUIRES_OK_ASYNC(ctx, + iterator->GetNext(IteratorContext(ctx), + &components, &end_of_sequence), + done); if (!end_of_sequence) { OP_REQUIRES_OK_ASYNC( |