Set environment variable to map the devices

parent 5bd29a11
......@@ -14,6 +14,7 @@ from tensorflow.core.framework import summary_pb2
import time
from bob.learn.tensorflow.datashuffler.OnlineSampling import OnLineSampling
os.environ["CUDA_VISIBLE_DEVICES"] = "1,2,3,0"
logger = bob.core.log.setup("bob.learn.tensorflow")
......@@ -276,7 +277,9 @@ class Trainer(object):
# Training
hdf5 = bob.io.base.HDF5File(os.path.join(self.temp_dir, 'model.hdf5'), 'w')
with tf.Session() as session:
config = tf.ConfigProto(log_device_placement=True)
config.gpu_options.allow_growth = True
with tf.Session(config=True) as session:
tf.initialize_all_variables().run()
......
Markdown is supported
0%
or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment