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

[XLA:Python] Don't attempt size-0 device-to-device transfers.

Browse Source 
PiperOrigin-RevId: 300767536
Change-Id: Id40cf08df386b39b447141a797aa59710d96e23a
This commit is contained in:
Peter Hawkins 2020-03-13 09:13:08 -07:00 committed by TensorFlower Gardener
parent 031a776366
commit a229588b9a
1 changed files with 5 additions and 3 deletions

8
tensorflow/compiler/xla/python/local_client.cc
View File

@ -597,9 +597,11 @@ StatusOr<std::unique_ptr<PyLocalBuffer>> PyLocalBuffer::CopyToDevice(
TF_RET_CHECK(input_buffer.size() == output_buffer.size())
<< "input: " << input_buffer.size()
<< " output: " << output_buffer.size();
TF_RETURN_IF_ERROR(transfer_local_device->ThenMemcpyDeviceToDevice(
transfer_stream, dst_local_device->compute_stream(), input_buffer,
output_buffer));
if (input_buffer.size() != 0) {
TF_RETURN_IF_ERROR(transfer_local_device->ThenMemcpyDeviceToDevice(
transfer_stream, dst_local_device->compute_stream(), input_buffer,
output_buffer));
}
}
// We hold on to the `src_device_buffer` until the transfer is finished.