eager: Remove unnecessary TFE_Context argument to TFE_OpSetDevice.
PiperOrigin-RevId: 168417999
This commit is contained in:
parent
86211d5545
commit
b0a397fcea
@ -289,11 +289,11 @@ static void TFE_OpSetDeviceHelper(TFE_Op* op, tensorflow::Device* device,
|
||||
}
|
||||
}
|
||||
|
||||
void TFE_OpSetDevice(TFE_Op* op, TFE_Context* ctx, const char* device_name,
|
||||
TF_Status* status) {
|
||||
void TFE_OpSetDevice(TFE_Op* op, const char* device_name, TF_Status* status) {
|
||||
tensorflow::Device* d = nullptr;
|
||||
if (device_name != nullptr && strlen(device_name) > 0) {
|
||||
status->status = ctx->session->device_mgr->LookupDevice(device_name, &d);
|
||||
status->status =
|
||||
op->ctx->session->device_mgr->LookupDevice(device_name, &d);
|
||||
if (!status->status.ok()) return;
|
||||
}
|
||||
TFE_OpSetDeviceHelper(op, d, status);
|
||||
|
@ -100,11 +100,8 @@ TF_CAPI_EXPORT extern TFE_Op* TFE_NewOp(TFE_Context* ctx, const char* op_or_func
|
||||
TF_Status* status);
|
||||
TF_CAPI_EXPORT extern void TFE_DeleteOp(TFE_Op* op);
|
||||
|
||||
// TODO(ashankar): TFE_OpSetDevice and TFE_Execute should not have a TFE_Context
|
||||
// parameter. Instead, the TFE_Context should be captured when creating the
|
||||
// TFE_Op.
|
||||
TF_CAPI_EXPORT extern void TFE_OpSetDevice(TFE_Op* op, TFE_Context* ctx,
|
||||
const char* device_name, TF_Status* status);
|
||||
TF_CAPI_EXPORT extern void TFE_OpSetDevice(TFE_Op* op, const char* device_name,
|
||||
TF_Status* status);
|
||||
|
||||
TF_CAPI_EXPORT extern void TFE_OpAddInput(TFE_Op* op, TFE_TensorHandle* h, TF_Status* status);
|
||||
|
||||
|
@ -297,7 +297,7 @@ void TFE_Py_Execute(TFE_Context* ctx, const char* device_name,
|
||||
TF_Status* out_status) {
|
||||
TFE_Op* op = TFE_NewOp(ctx, op_name, out_status);
|
||||
if (TF_GetCode(out_status) != TF_OK) return;
|
||||
TFE_OpSetDevice(op, ctx, device_name, out_status);
|
||||
TFE_OpSetDevice(op, device_name, out_status);
|
||||
if (TF_GetCode(out_status) == TF_OK) {
|
||||
for (int i = 0; i < inputs->size() && TF_GetCode(out_status) == TF_OK;
|
||||
++i) {
|
||||
|
Loading…
Reference in New Issue
Block a user