From 9ec3608eca34aaa17c3d7e7c41afc940043db94c Mon Sep 17 00:00:00 2001 From: ustcfriend Date: Wed, 13 Sep 2017 13:36:06 +0800 Subject: [PATCH] Fix resnet crash by setting config.gpu_options.allow_growth = True. (#971) --- examples/resnet/resnet_main.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/examples/resnet/resnet_main.py b/examples/resnet/resnet_main.py index 51bc4f761..a8f6d711e 100644 --- a/examples/resnet/resnet_main.py +++ b/examples/resnet/resnet_main.py @@ -97,6 +97,7 @@ class ResNetTrainActor(object): self.model = resnet_model.ResNet(hps, images, labels, "train") self.model.build_graph() config = tf.ConfigProto(allow_soft_placement=True) + config.gpu_options.allow_growth = True sess = tf.Session(config=config) self.model.variables.set_session(sess) self.coord = tf.train.Coordinator() @@ -144,6 +145,7 @@ class ResNetTestActor(object): self.model = resnet_model.ResNet(hps, images, labels, "eval") self.model.build_graph() config = tf.ConfigProto(allow_soft_placement=True) + config.gpu_options.allow_growth = True sess = tf.Session(config=config) self.model.variables.set_session(sess) self.coord = tf.train.Coordinator()