diff options
Diffstat (limited to 'tensorflow/core/kernels/function_ops.cc')
-rw-r--r-- | tensorflow/core/kernels/function_ops.cc | 12 |
1 files changed, 5 insertions, 7 deletions
diff --git a/tensorflow/core/kernels/function_ops.cc b/tensorflow/core/kernels/function_ops.cc index f469f41e06..facac10f66 100644 --- a/tensorflow/core/kernels/function_ops.cc +++ b/tensorflow/core/kernels/function_ops.cc @@ -296,21 +296,19 @@ class RemoteCallOp : public AsyncOpKernel { void ComputeAsync(OpKernelContext* ctx, DoneCallback done) override { const Tensor* target; OP_REQUIRES_OK_ASYNC(ctx, ctx->input("target", &target), done); - AttrValueMap attr_values = func_.attr(); - AttrValue v; const string& target_device = DeviceNameUtils::CanonicalizeDeviceName(target->scalar<string>()()); - v.set_s(target_device); - AddAttr("_target", v, &attr_values); FunctionLibraryRuntime* lib = ctx->function_library(); OP_REQUIRES_ASYNC(ctx, lib != nullptr, errors::Internal("No function library is provided."), done); + AttrValueMap attr_values = func_.attr(); FunctionLibraryRuntime::Handle handle; - OP_REQUIRES_OK_ASYNC( - ctx, lib->Instantiate(func_.name(), AttrSlice(&attr_values), &handle), - done); + OP_REQUIRES_OK_ASYNC(ctx, + lib->Instantiate(func_.name(), AttrSlice(&attr_values), + {target_device}, &handle), + done); OpInputList arguments; OP_REQUIRES_OK_ASYNC(ctx, ctx->input_list("args", &arguments), done); |