//    out the same buffers until the transfer has completed.
    status = host_to_device_stream_->DoHostCallback([ref]() { ref.Unref(); });
    done(status);
  } else {
    status = host_to_device_stream_->DoHostCallback([ref, done]() {
      ref.Unref();
      done(absl::OkStatus());
    });
    if (!status.ok()) {
      done(status);
    }
  }
}

    // might not be enqueued to the stream yet, we put it on destination stream.
    TensorReference input_reference(*input);
    std::move(return_substream).release();
    return dst_device_to_device_stream->DoHostCallback(
        [input_reference, done = std::move(done),
         device_to_device_master_stream, dst_device_to_device_stream] {
          if (device_to_device_master_stream) {