diff options
Diffstat (limited to 'tensorflow/stream_executor/stream.cc')
-rw-r--r-- | tensorflow/stream_executor/stream.cc | 13 |
1 files changed, 13 insertions, 0 deletions
diff --git a/tensorflow/stream_executor/stream.cc b/tensorflow/stream_executor/stream.cc index a42a469df5..9efd34de24 100644 --- a/tensorflow/stream_executor/stream.cc +++ b/tensorflow/stream_executor/stream.cc @@ -5294,6 +5294,19 @@ Stream &Stream::ThenDoHostCallback(std::function<void()> callback) { return *this; } +Stream &Stream::ThenDoHostCallbackWithStatus( + std::function<port::Status()> callback) { + VLOG_CALL(PARAM(callback)); + + if (ok()) { + CheckError(parent_->HostCallback(this, std::move(callback))); + } else { + LOG(WARNING) << "stream " << DebugStreamPointers() + << " was in error state before adding host callback"; + } + return *this; +} + Stream &Stream::ThenFft(fft::Plan *plan, const DeviceMemory<std::complex<float>> &input, DeviceMemory<std::complex<float>> *output) { |