diff options
| author |  Alexandre Passos <apassos@google.com> | 2018-03-01 09:27:57 -0800 |
|---|---|---|
| committer |  TensorFlower Gardener <gardener@tensorflow.org> | 2018-03-01 09:31:56 -0800 |
| commit | c65343d282cdf5ccf4f7d3229f6c492fec344f8d (patch) | |
| tree | 36573dd621786a455af99874a4798a99af93324d /tensorflow/python/lib | |
| parent | 03de984caa1f1403d4417357b67e96dfb7edbc3e (diff) | |
Keep track of eager op device for tensor handles. Force-colocates ops using resources with the resources.
PiperOrigin-RevId: 187488175
Diffstat (limited to 'tensorflow/python/lib')
| -rw-r--r-- | tensorflow/python/lib/core/py_func.cc | 5 |
1 files changed, 3 insertions, 2 deletions
diff --git a/tensorflow/python/lib/core/py_func.cc b/tensorflow/python/lib/core/py_func.cc index e0422ef80a..343415b264 100644 --- a/tensorflow/python/lib/core/py_func.cc +++ b/tensorflow/python/lib/core/py_func.cc @@ -79,10 +79,11 @@ Status MakeArgTuple(const PyCall* call, PyObject** tuple) { const Tensor& t = call->ins[i]; if (call->eager) { if (call->gpu) { - arg = EagerTensorFromHandle(new TFE_TensorHandle(t, call->device)); + arg = EagerTensorFromHandle( + new TFE_TensorHandle(t, call->device, call->device)); } else { // TFE_TensorHandle assumes that CPU is identified by `nullptr`. - arg = EagerTensorFromHandle(new TFE_TensorHandle(t, nullptr)); + arg = EagerTensorFromHandle(new TFE_TensorHandle(t, nullptr, nullptr)); } if (arg == nullptr) { return errors::Internal("Unable to procure EagerTensor from Tensor."); |