diff options
| author |  Alexandre Passos <apassos@google.com> | 2018-02-23 17:19:00 -0800 |
|---|---|---|
| committer |  TensorFlower Gardener <gardener@tensorflow.org> | 2018-02-23 17:23:07 -0800 |
| commit | b1cc57604cadb4251efeb764074c9138d4e24521 (patch) | |
| tree | 25d720269aeff317c124d5a0c80d767f16d52408 | |
| parent | ce4ae5bed9b47f49b085d9d8287cee2fcc5d42ac (diff) | |
Dropped from previous change.
PiperOrigin-RevId: 186846681
| -rw-r--r-- | tensorflow/c/eager/c_api.cc | 8 |
1 files changed, 3 insertions, 5 deletions
diff --git a/tensorflow/c/eager/c_api.cc b/tensorflow/c/eager/c_api.cc index f615e3f11d..c27a7129fa 100644 --- a/tensorflow/c/eager/c_api.cc +++ b/tensorflow/c/eager/c_api.cc @@ -303,11 +303,9 @@ void TFE_OpSetXLACompilation(TFE_Op* op, unsigned char enable) { void TFE_OpAddInput(TFE_Op* op, TFE_TensorHandle* h, TF_Status* status) { // Questionable heuristic ... - // - // Motivation: After an 'op' is placed on GPU because some of its earlier - // inputs are on GPU, we want to keep the 'op' there, even if some later - // inputs of it are not on GPU. - if (IsCPU(op->device) && !IsCPU(h->d)) { + // - If a device was explicitly set on the op, always use that. + // - If not, place on the first non-host device seen. + if (op->device == nullptr && !IsCPU(h->d)) { op->device = h->d; } if (!status->status.ok()) return; |