diff options
Diffstat (limited to 'tensorflow/core/kernels/data/captured_function.cc')
| -rw-r--r-- | tensorflow/core/kernels/data/captured_function.cc | 12 |
1 files changed, 12 insertions, 0 deletions
diff --git a/tensorflow/core/kernels/data/captured_function.cc b/tensorflow/core/kernels/data/captured_function.cc index ee58341cfd..82da385405 100644 --- a/tensorflow/core/kernels/data/captured_function.cc +++ b/tensorflow/core/kernels/data/captured_function.cc @@ -214,6 +214,9 @@ Status CapturedFunction::Run(IteratorContext* ctx, std::vector<Tensor>&& args, }); f_opts.step_container = &step_container; f_opts.runner = ctx->runner(); + if (ctx->lib()->device()->device_type() != DEVICE_CPU) { + f_opts.create_rendezvous = true; + } // TODO(mrry): Add cancellation manager support to IteratorContext // so that we can cancel running map functions. The local // cancellation manager here is created so that we can run kernels @@ -248,6 +251,9 @@ Status CapturedFunction::RunWithBorrowedArgs(IteratorContext* ctx, }); f_opts.step_container = &step_container; f_opts.runner = ctx->runner(); + if (ctx->lib()->device()->device_type() != DEVICE_CPU) { + f_opts.create_rendezvous = true; + } // TODO(mrry): Add cancellation manager support to IteratorContext // so that we can cancel running map functions. The local // cancellation manager here is created so that we can run kernels @@ -304,6 +310,9 @@ Status CapturedFunction::RunInstantiated(const std::vector<Tensor>& args, }); f_opts.step_container = &step_container; f_opts.runner = runner; + if (lib->device()->device_type() != DEVICE_CPU) { + f_opts.create_rendezvous = true; + } // TODO(mrry): Add cancellation manager support to IteratorContext // so that we can cancel running map functions. The local // cancellation manager here is created so that we can run kernels @@ -351,6 +360,9 @@ void CapturedFunction::RunAsync(IteratorContext* ctx, }); f_opts.step_container = step_container; f_opts.runner = ctx->runner(); + if (ctx->lib()->device()->device_type() != DEVICE_CPU) { + f_opts.create_rendezvous = true; + } // TODO(mrry): Add cancellation manager support to IteratorContext // so that we can cancel running map functions. The local // cancellation manager here is created so that we can run kernels |