From f12c967893eaab978963d9bb887bff7bcd099839 Mon Sep 17 00:00:00 2001 From: Benoit Steiner Date: Thu, 20 Apr 2017 08:27:45 -0800 Subject: [PATCH] Create a per session inter op thread pool instead of relying on a global thread pool: this ensures that the thread pool is reset whenever the session is reset. Change: 153718067 --- tensorflow/core/grappler/clusters/single_machine.cc | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/tensorflow/core/grappler/clusters/single_machine.cc b/tensorflow/core/grappler/clusters/single_machine.cc index 9e96cc1e912..219d95acdfd 100644 --- a/tensorflow/core/grappler/clusters/single_machine.cc +++ b/tensorflow/core/grappler/clusters/single_machine.cc @@ -41,7 +41,10 @@ SingleMachine::SingleMachine(int timeout_s, int num_cpu_cores, int num_gpus) } CHECK_GE(num_cpu_cores, 1); options_.config.set_intra_op_parallelism_threads(num_cpu_cores); - options_.config.set_inter_op_parallelism_threads(num_cpu_cores); + // Create a session specific thread pool to ensure the threads are reset when + // the session is reset. + options_.config.add_session_inter_op_thread_pool()->set_num_threads( + num_cpu_cores); if (timeout_s > 0) { options_.config.set_operation_timeout_in_ms(timeout_s * 1000); }