From d2a6f7958a5126c6234e3f47eb80578deb2fc5a4 Mon Sep 17 00:00:00 2001 From: Richard Liaw Date: Fri, 23 Aug 2019 18:29:40 -0700 Subject: [PATCH] [tune] Fix for keras threading (#5517) --- python/ray/tune/examples/utils.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/python/ray/tune/examples/utils.py b/python/ray/tune/examples/utils.py index f40707a01..bee27a889 100644 --- a/python/ray/tune/examples/utils.py +++ b/python/ray/tune/examples/utils.py @@ -2,6 +2,7 @@ from __future__ import absolute_import from __future__ import division from __future__ import print_function +import tensorflow as tf import keras from keras.datasets import mnist from keras import backend as K @@ -52,8 +53,8 @@ def set_keras_threads(threads): # We set threads here to avoid contention, as Keras # is heavily parallelized across multiple cores. K.set_session( - K.tf.Session( - config=K.tf.ConfigProto( + tf.Session( + config=tf.ConfigProto( intra_op_parallelism_threads=threads, inter_op_parallelism_threads=threads)))