experiments/STT-tensorflow
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

Enable collective ops in defun

Browse Source 
PiperOrigin-RevId: 226942243
This commit is contained in:
Yuefeng Zhou 2018-12-26 12:17:17 -08:00 committed by TensorFlower Gardener
parent 48f2893491
commit 742778f04e
1 changed files with 1 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

1
tensorflow/core/kernels/partitioned_function_ops.cc
View File

@ -439,6 +439,7 @@ class PartitionedCallOp : public AsyncOpKernel {
opts.step_container = ctx->step_container(); opts.step_container = ctx->step_container();
opts.cancellation_manager = ctx->cancellation_manager(); opts.cancellation_manager = ctx->cancellation_manager();
opts.stats_collector = ctx->stats_collector(); opts.stats_collector = ctx->stats_collector();
opts.collective_executor = ctx->collective_executor();
// TODO(akshayka): Consider selecting a runner on a per-device basis, i.e., // TODO(akshayka): Consider selecting a runner on a per-device basis, i.e.,
// using device-specific threadpools when available. // using device-specific threadpools when available.
opts.runner = ctx->runner(); opts.runner = ctx->runner();