Enable collective ops in defun
PiperOrigin-RevId: 226942243
This commit is contained in:
parent
48f2893491
commit
742778f04e
@ -439,6 +439,7 @@ class PartitionedCallOp : public AsyncOpKernel {
|
||||
opts.step_container = ctx->step_container();
|
||||
opts.cancellation_manager = ctx->cancellation_manager();
|
||||
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.,
|
||||
// using device-specific threadpools when available.
|
||||
opts.runner = ctx->runner();
|
||||
|
Loading…
Reference in New Issue
Block a user