From f7a23a4459adbd0ae1de2564aff1895872798d94 Mon Sep 17 00:00:00 2001 From: Falcon Dai Date: Wed, 26 Jul 2017 15:16:02 -0500 Subject: [PATCH] fix #18 run on GPU --- .gitignore | 2 ++ run_summarization.py | 3 ++- 2 files changed, 4 insertions(+), 1 deletion(-) create mode 100644 .gitignore diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..fe2d248 --- /dev/null +++ b/.gitignore @@ -0,0 +1,2 @@ +*.pyc +*logs/ diff --git a/run_summarization.py b/run_summarization.py index 8639af9..c4673b9 100644 --- a/run_summarization.py +++ b/run_summarization.py @@ -122,7 +122,8 @@ def setup_training(model, batcher): train_dir = os.path.join(FLAGS.log_root, "train") if not os.path.exists(train_dir): os.makedirs(train_dir) - default_device = tf.device('/cpu:0') + # use the first available GPU + default_device = tf.device('/gpu:0') with default_device: model.build_graph() # build the graph if FLAGS.convert_to_coverage_model: