session_config

lefnire · Feb 10, 2018 · 6f8f8e9 · 6f8f8e9
1 parent aa793c2
commit 6f8f8e9
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/tensorforce/models/model.py b/tensorforce/models/model.py
@@ -533,13 +533,14 @@ def init_fn(scaffold, session):
         # tf.train.NanTensorHook(loss_tensor, fail_on_nan_loss=True)
         # tf.train.ProfilerHook(save_steps=None, save_secs=None, output_dir='', show_dataflow=True, show_memory=False)
 
+        session_config = tf.ConfigProto(gpu_options=tf.GPUOptions(per_process_gpu_memory_fraction=.9/2))
         if self.distributed_spec is None:
             # TensorFlow non-distributed monitored session object
             self.monitored_session = tf.train.SingularMonitoredSession(
                 hooks=hooks,
                 scaffold=self.scaffold,
                 master='',  # Default value.
-                config=None,  # self.distributed_spec.get('session_config'),
+                config=session_config,
                 checkpoint_dir=None
             )