Implement DoHostCallbackWithStatus to allow callbacks to return a status

PiperOrigin-RevId: 207714420

1 files changed, 5 insertions, 0 deletions

diff --git a/tensorflow/stream_executor/stream_executor_pimpl.cc b/tensorflow/stream_executor/stream_executor_pimpl.cc
index 2e0137a485..9515d8e62a 100644
--- a/tensorflow/stream_executor/stream_executor_pimpl.cc
+++ b/tensorflow/stream_executor/stream_executor_pimpl.cc
@@ -699,6 +699,11 @@ bool StreamExecutor::HostCallback(Stream *stream,
   return implementation_->HostCallback(stream, std::move(callback));
 }
 
+bool StreamExecutor::HostCallback(Stream *stream,
+                                  std::function<port::Status()> callback) {
+  return implementation_->HostCallback(stream, std::move(callback));
+}
+
 port::Status StreamExecutor::AllocateEvent(Event *event) {
   return implementation_->AllocateEvent(event);
 }