experiments/STT-tensorflow
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

Support packed TensorHandles on CustomDevices.

Browse Source 
PiperOrigin-RevId: 324851163
Change-Id: I017665b5e2bd90da37f4d45aa94e1278be7bd2f7
This commit is contained in:
Yujing Zhang 2020-08-04 11:17:45 -07:00 committed by TensorFlower Gardener
parent 7a2383c18f
commit 4e03f13e6a
1 changed files with 1 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

3
tensorflow/core/common_runtime/eager/tensor_handle.cc
View File

@ -316,8 +316,7 @@ Status TensorHandle::CreatePackedHandle(std::vector<TensorHandle*>&& handles,
std::vector<string> devices;
for (auto* handle : handles) {
if (VariantDeviceIsCustom(handle->device())) {
return errors::InvalidArgument(
"CustomDevice is not supported for packing.");
devices.push_back(absl::get<CustomDevice*>(handle->device())->name());
} else {
devices.push_back(handle->op_device() ? handle->op_device()->name()
: ctx->HostCPU()->name());