eager: Remove unnecessary TFE_Context argument to TFE_OpSetDevice.

PiperOrigin-RevId: 168417999
This commit is contained in:
Asim Shankar 2017-09-12 12:16:59 -07:00 committed by TensorFlower Gardener
parent 86211d5545
commit b0a397fcea
3 changed files with 6 additions and 9 deletions

View File

@ -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);

View File

@ -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);

View File

@ -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) {