Do not assert for null device memory pointer if size is zero.

PiperOrigin-RevId: 307658410
Change-Id: I198173e9a3e4d967075f9b31da211f6c73c5572f
This commit is contained in:
Davide Libenzi 2020-04-21 12:29:40 -07:00 committed by TensorFlower Gardener
parent 0fa30f59f3
commit d914e48cd1
1 changed files with 4 additions and 2 deletions

View File

@ -588,7 +588,8 @@ xla::StatusOr<xla::ShapedBuffer> XRTTupleAllocation::ToShapedBuffer() {
allocator_->platform(), device_ordinal_); allocator_->platform(), device_ordinal_);
for (const auto& index_buffer : buffers_) { for (const auto& index_buffer : buffers_) {
if (index_buffer.second == nullptr || if (index_buffer.second == nullptr ||
index_buffer.second->allocation().is_null()) { (index_buffer.second->allocation().is_null() &&
index_buffer.second->allocation().size() > 0)) {
return errors::InvalidArgument("Literal buffer at index ", return errors::InvalidArgument("Literal buffer at index ",
index_buffer.first.ToString(), index_buffer.first.ToString(),
" has been released"); " has been released");
@ -652,7 +653,8 @@ xla::StatusOr<xla::ExecutionInput> XRTTupleAllocation::ToExecutionInput(
xla::ExecutionInput result(on_device_shape()); xla::ExecutionInput result(on_device_shape());
for (const auto& index_buffer : buffers_) { for (const auto& index_buffer : buffers_) {
if (index_buffer.second == nullptr || if (index_buffer.second == nullptr ||
index_buffer.second->allocation().is_null()) { (index_buffer.second->allocation().is_null() &&
index_buffer.second->allocation().size() > 0)) {
return errors::InvalidArgument("Literal buffer at index ", return errors::InvalidArgument("Literal buffer at index ",
index_buffer.first.ToString(), index_buffer.first.ToString(),
" has been released"); " has been released");