diff options
author | 2018-07-24 13:50:22 -0700 | |
---|---|---|
committer | 2018-07-24 13:55:16 -0700 | |
commit | ee0bd6ef450b388fadea63b31b65b13bd12f17d6 (patch) | |
tree | a7575f0449f0235a8cf9a05a3207f6f60990e943 /tensorflow/compiler/xla/service/hlo_runner.cc | |
parent | 6b2ae9b3da572d2f8e1eccdf2922efce43cbecd4 (diff) |
Automated rollback of commit 0ea6847c892497afdd20c1150fee1e532612ca17
PiperOrigin-RevId: 205885304
Diffstat (limited to 'tensorflow/compiler/xla/service/hlo_runner.cc')
-rw-r--r-- | tensorflow/compiler/xla/service/hlo_runner.cc | 9 |
1 files changed, 7 insertions, 2 deletions
diff --git a/tensorflow/compiler/xla/service/hlo_runner.cc b/tensorflow/compiler/xla/service/hlo_runner.cc index 4f0569f405..b2725e2918 100644 --- a/tensorflow/compiler/xla/service/hlo_runner.cc +++ b/tensorflow/compiler/xla/service/hlo_runner.cc @@ -180,8 +180,12 @@ StatusOr<ScopedShapedBuffer> HloRunner::ExecuteWithDeviceBuffers( TF_ASSIGN_OR_RETURN(std::unique_ptr<Executable> executable, CreateExecutable(std::move(module), run_hlo_passes)); - return executable->ExecuteOnStreamWrapper(&service_run_options, - /*profile=*/profile, arguments); + TF_ASSIGN_OR_RETURN( + ScopedShapedBuffer retval, + executable->ExecuteOnStreamWrapper(&service_run_options, + /*profile=*/profile, arguments)); + TF_RETURN_IF_ERROR(stream.BlockHostUntilDone()); + return std::move(retval); } StatusOr<ScopedShapedBuffer> HloRunner::ExecuteWithDeviceBuffers( @@ -309,6 +313,7 @@ StatusOr<std::vector<std::unique_ptr<Literal>>> HloRunner::ExecuteReplicated( std::vector<std::unique_ptr<Literal>> exec_results; for (int64 i = 0; i < options.num_replicas; ++i) { + TF_RETURN_IF_ERROR(streams[i]->BlockHostUntilDone()); TF_ASSIGN_OR_RETURN(std::unique_ptr<Literal> literal, backend().transfer_manager()->TransferLiteralFromDevice( streams[i].get(), results[i])); |