From af8772fa69fb5593ec4d3e297602723705c922e2 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Tue, 9 Jul 2019 10:03:22 -0700
Subject: [PATCH] [XLA:Python] Fix potential issue in CopyToDevice where source
 device buffer might get deallocated before transfer completes.

PiperOrigin-RevId: 257216762
---
 tensorflow/compiler/xla/python/local_client.cc | 4 ++++
 1 file changed, 4 insertions(+)

diff --git a/tensorflow/compiler/xla/python/local_client.cc b/tensorflow/compiler/xla/python/local_client.cc
index d7f27d00f76..2057788c12d 100644
--- a/tensorflow/compiler/xla/python/local_client.cc
+++ b/tensorflow/compiler/xla/python/local_client.cc
@@ -652,6 +652,10 @@ StatusOr<std::unique_ptr<PyLocalBuffer>> PyLocalBuffer::CopyToDevice(
         output_buffer));
   }
 
+  // We hold on to the `src_device_buffer` until the transfer is finished.
+  src_device.ThenRelease(src_device_to_device_stream,
+                         std::move(src_device_buffer));
+
   // Write new tuple buffers. The destination buffers have different addresses,
   // so we must construct tuple buffers from scratch instead of copying them.
   if (dst_buffer.on_device_shape().IsTuple()) {