diff options
| author |  Rohan Jain <rohanj@google.com> | 2018-07-25 16:50:39 -0700 |
|---|---|---|
| committer |  TensorFlower Gardener <gardener@tensorflow.org> | 2018-07-25 16:53:32 -0700 |
| commit | 773ee806ceb1feab77ac02f424814eb4ab503170 (patch) | |
| tree | eb12eb67b2d2786a6621c997292493dac13e9ad7 /tensorflow | |
| parent | 1e5854fb157a6be7a6a2b4bad97bce50c264e678 (diff) | |
Fixing bug in MultiDeviceIterator where we were using OP_REQUIRES_OK in an AsyncOpKernel. Replacing it with OP_REQUIRES_OK_ASYNC.
PiperOrigin-RevId: 206078496
Diffstat (limited to 'tensorflow')
| -rw-r--r-- | tensorflow/contrib/data/kernels/prefetching_kernels.cc | 9 |
1 files changed, 5 insertions, 4 deletions
diff --git a/tensorflow/contrib/data/kernels/prefetching_kernels.cc b/tensorflow/contrib/data/kernels/prefetching_kernels.cc index 6edc61b2c2..32f03ca683 100644 --- a/tensorflow/contrib/data/kernels/prefetching_kernels.cc +++ b/tensorflow/contrib/data/kernels/prefetching_kernels.cc @@ -791,16 +791,17 @@ class MultiDeviceIteratorGetNextFromShardOp : public AsyncOpKernel { void ComputeAsync(OpKernelContext* ctx, DoneCallback done) override { const Tensor* tensor_shard_num; - OP_REQUIRES_OK(ctx, ctx->input("shard_num", &tensor_shard_num)); + OP_REQUIRES_OK_ASYNC(ctx, ctx->input("shard_num", &tensor_shard_num), done); int32 shard_num = tensor_shard_num->scalar<int32>()(); const Tensor* tensor_incarnation_id; - OP_REQUIRES_OK(ctx, ctx->input("incarnation_id", &tensor_incarnation_id)); + OP_REQUIRES_OK_ASYNC( + ctx, ctx->input("incarnation_id", &tensor_incarnation_id), done); int64 incarnation_id = tensor_incarnation_id->scalar<int64>()(); MultiDeviceIterator* iterator; - OP_REQUIRES_OK(ctx, - LookupResource(ctx, HandleFromInput(ctx, 0), &iterator)); + OP_REQUIRES_OK_ASYNC( + ctx, LookupResource(ctx, HandleFromInput(ctx, 0), &iterator), done); thread_pool_->Schedule(std::bind( [ctx, iterator, shard_num, incarnation_id](DoneCallback done) { std::vector<Tensor> components; |