From a7d75e5df3d1f8c0c14b22dafa1c651fcf4e26e6 Mon Sep 17 00:00:00 2001 From: Arthur Juliani Date: Thu, 6 Sep 2018 13:09:14 -0700 Subject: [PATCH] Add fix for multiple instances on a single GPU --- ml-agents/mlagents/trainers/trainer_controller.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/ml-agents/mlagents/trainers/trainer_controller.py b/ml-agents/mlagents/trainers/trainer_controller.py index 06f8736aac..be53dcee89 100644 --- a/ml-agents/mlagents/trainers/trainer_controller.py +++ b/ml-agents/mlagents/trainers/trainer_controller.py @@ -291,7 +291,10 @@ def start_learning(self): tf.reset_default_graph() - with tf.Session() as sess: + # Prevent a single session from taking all GPU memory. + config = tf.ConfigProto() + config.gpu_options.allow_growth = True + with tf.Session(config=config) as sess: self._initialize_trainers(trainer_config, sess) for _, t in self.trainers.items(): self.logger.info(t)