From 64acbc93ad67a1dd2869fa4b5cb5ae8a04ac3b34 Mon Sep 17 00:00:00 2001
From: Tiago Freitas Pereira <tiagofrepereira@gmail.com>
Date: Sun, 23 Oct 2016 18:14:27 +0200
Subject: [PATCH] Set environment variable to map the devices

---
 bob/learn/tensorflow/trainers/Trainer.py | 5 ++++-
 1 file changed, 4 insertions(+), 1 deletion(-)

diff --git a/bob/learn/tensorflow/trainers/Trainer.py b/bob/learn/tensorflow/trainers/Trainer.py
index d3500ab6..ecd30aa1 100644
--- a/bob/learn/tensorflow/trainers/Trainer.py
+++ b/bob/learn/tensorflow/trainers/Trainer.py
@@ -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()
 
-- 
GitLab