Allow TFE_TensorHandleCopyToDevice to have the same device as src and

destination. It will reuse the same underlying buffer in those cases.

PiperOrigin-RevId: 164909906

1 files changed, 10 insertions, 2 deletions

diff --git a/tensorflow/c/eager/c_api_test.cc b/tensorflow/c/eager/c_api_test.cc
index 797506422b..6614df78d9 100644
--- a/tensorflow/c/eager/c_api_test.cc
+++ b/tensorflow/c/eager/c_api_test.cc
@@ -170,14 +170,22 @@ TEST(CAPI, TensorHandleCopyBetweenDevices) {
       ADD_FAILURE() << tag << " -- " << TF_Message(status.get());
       continue;
     }
+    // Copy from device to the same device.
+    TFE_TensorHandle* hdevice2 =
+        TFE_TensorHandleCopyToDevice(hdevice, ctx, name.c_str(), status.get());
+    if (TF_GetCode(status.get()) != TF_OK) {
+      ADD_FAILURE() << tag << " -- " << TF_Message(status.get());
+      continue;
+    }
+    TFE_DeleteTensorHandle(hdevice);
     // Copy back to CPU
     TFE_TensorHandle* hcopy =
-        TFE_TensorHandleCopyToDevice(hdevice, ctx, kCPUDevice, status.get());
+        TFE_TensorHandleCopyToDevice(hdevice2, ctx, kCPUDevice, status.get());
     if (TF_GetCode(status.get()) != TF_OK) {
       ADD_FAILURE() << tag << " -- " << TF_Message(status.get());
       continue;
     }
-    TFE_DeleteTensorHandle(hdevice);
+    TFE_DeleteTensorHandle(hdevice2);
 
     // Ensure that the contents are the same!
     TF_Tensor* tcopy = TFE_TensorHandleResolve(hcopy, status.get());