diff options
Diffstat (limited to 'tensorflow/c/eager/c_api.cc')
| -rw-r--r-- | tensorflow/c/eager/c_api.cc | 3 |
1 files changed, 2 insertions, 1 deletions
diff --git a/tensorflow/c/eager/c_api.cc b/tensorflow/c/eager/c_api.cc index e4e33016c2..ecba8c0109 100644 --- a/tensorflow/c/eager/c_api.cc +++ b/tensorflow/c/eager/c_api.cc @@ -190,6 +190,7 @@ TFE_TensorHandle* TFE_TensorHandleCopyToDevice(TFE_TensorHandle* h, bool is_same_device = (srcd == dstd) || (DeviceName(srcd) == DeviceName(dstd)); const bool dst_cpu = IsCPU(dstd); + const bool src_cpu = IsCPU(srcd); if (is_same_device) { return new TFE_TensorHandle(h->t, dst_cpu ? nullptr : dstd); } @@ -213,7 +214,7 @@ TFE_TensorHandle* TFE_TensorHandleCopyToDevice(TFE_TensorHandle* h, return new TFE_TensorHandle(dst, dst_cpu ? nullptr : dstd); } tensorflow::DeviceContext* src_device_context = nullptr; - if (!IsCPU(srcd)) { + if (!src_cpu) { src_device_context = srcd->tensorflow_gpu_device_info()->default_context; } tensorflow::DeviceContext* dst_device_context = nullptr; |