wassname/ray
1
0
Fork 0
mirror of https://github.com/wassname/ray.git synced 2026-07-06 04:44:08 +08:00
Code Issues Packages Projects Releases Wiki Activity

Fix resnet crash by setting config.gpu_options.allow_growth = True. (#971)

Browse Source
This commit is contained in:
ustcfriend
2017-09-13 13:36:06 +08:00
committed by Robert Nishihara
parent 9f42ef6a4f
commit 9ec3608eca
1 changed files with 2 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
examples/resnet/resnet_main.py
+2
View File
@@ -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()