diff options
Diffstat (limited to 'tensorflow/core/common_runtime/direct_session.cc')
| -rw-r--r-- | tensorflow/core/common_runtime/direct_session.cc | 8 |
1 files changed, 4 insertions, 4 deletions
diff --git a/tensorflow/core/common_runtime/direct_session.cc b/tensorflow/core/common_runtime/direct_session.cc index 0695278c0d..bf1d78ec65 100644 --- a/tensorflow/core/common_runtime/direct_session.cc +++ b/tensorflow/core/common_runtime/direct_session.cc @@ -602,7 +602,7 @@ Status DirectSession::RunInternal(int64 step_id, const RunOptions& run_options, if (tracer) { TF_RETURN_IF_ERROR(tracer->Stop()); - TF_RETURN_IF_ERROR(tracer->Collect(args.stats_collector)); + TF_RETURN_IF_ERROR(tracer->Collect(run_state.collector.get())); } { @@ -618,8 +618,8 @@ Status DirectSession::RunInternal(int64 step_id, const RunOptions& run_options, &session_state_)); } - if (args.stats_collector) { - args.stats_collector->Finalize(); + if (run_state.collector) { + run_state.collector->Finalize(); } // Build and return the cost model as instructed. @@ -634,7 +634,7 @@ Status DirectSession::RunInternal(int64 step_id, const RunOptions& run_options, } mutex_lock l(executor_lock_); - args.stats_collector->BuildCostModel(&cost_model_manager_, device_to_graph); + run_state.collector->BuildCostModel(&cost_model_manager_, device_to_graph); // annotate stats onto cost graph. CostGraphDef* cost_graph = run_metadata->mutable_cost_graph(); |