diff options
| -rw-r--r-- | tensorflow/core/distributed_runtime/master_session.cc | 8 |
1 files changed, 6 insertions, 2 deletions
diff --git a/tensorflow/core/distributed_runtime/master_session.cc b/tensorflow/core/distributed_runtime/master_session.cc index 9231fae362..870970b7ca 100644 --- a/tensorflow/core/distributed_runtime/master_session.cc +++ b/tensorflow/core/distributed_runtime/master_session.cc @@ -569,11 +569,15 @@ Status MasterSession::ReffedClientGraph::RunPartitions( bool success = cm->RegisterCallback(token, [&calls]() { calls.StartCancel(); }); if (!success) { - return errors::Cancelled("Step was cancelled"); + calls.StartCancel(); } calls.Wait(); - cm->DeregisterCallback(token); call_opts->ClearCancelCallback(); + if (success) { + cm->DeregisterCallback(token); + } else { + return errors::Cancelled("Step was cancelled"); + } // Collects fetches. Status status = calls.status(); |