From 51e4dfe363c55fdcb05fd5812a2d60db4459cf00 Mon Sep 17 00:00:00 2001
From: Chen Li <chencli@google.com>
Date: Wed, 2 Oct 2024 11:31:19 -0700
Subject: [PATCH] Internal change: fix non-copyable object from
 TF_ASSIGN_OR_RETURN

PiperOrigin-RevId: 681529749
---
 xla/pjrt/pjrt_stream_executor_client.cc | 10 ++++++----
 1 file changed, 6 insertions(+), 4 deletions(-)

diff --git a/xla/pjrt/pjrt_stream_executor_client.cc b/xla/pjrt/pjrt_stream_executor_client.cc
index cb92e535e10f6..34e856e79f2b4 100644
--- a/xla/pjrt/pjrt_stream_executor_client.cc
+++ b/xla/pjrt/pjrt_stream_executor_client.cc
@@ -2465,13 +2465,15 @@ PjRtStreamExecutorLoadedExecutable::MakeExecutionInputsAndWaitForEvents(
       client_->client()->backend().transfer_manager();
   // Lift tuple_handle outside the conditional so that the event it returns is
   // not destroyed until after the loop below that waits on events.
-  std::optional<TupleHandle> tuple_handle;
+  absl::StatusOr<TupleHandle> tuple_handle;
   if (parameter_is_tupled_arguments_ && !options.arguments_are_tupled) {
-    TF_ASSIGN_OR_RETURN(
-        tuple_handle,
+    tuple_handle =
         MakeTupleHelper(client_, device_state, options.strict_shape_checking,
                         executable_parameter_shapes[0], argument_handles,
-                        device_buffers, device_ordinal));
+                        device_buffers, device_ordinal);
+    if (!tuple_handle.ok()) {
+      return tuple_handle.status();
+    }
     events.insert(tuple_handle->event.get());
     execution_inputs.emplace_back(std::move(tuple_handle->execution_input));
   } else {