Allow using a single GPU

dab1517e · kaiix · 761c59c4 · dab1517e
Commit dab1517e authored Aug 31, 2016 by kaiix
Hide whitespace changes
Inline Side-by-side

Showing with 2 additions and 1 deletion

textsum/seq2seq_attention_model.py textsum/seq2seq_attention_model.py +2 -1

No files found.
--- a/textsum/seq2seq_attention_model.py
+++ b/textsum/seq2seq_attention_model.py
@@ -105,7 +105,8 @@ class Seq2SeqAttentionModel(object):
    if self._num_gpus == 0:
      return ''
    dev = '/gpu:%d' % self._cur_gpu
-    self._cur_gpu = (self._cur_gpu + 1) % (self._num_gpus-1)
+    if self._num_gpus > 1:
+      self._cur_gpu = (self._cur_gpu + 1) % (self._num_gpus-1)
    return dev

  def _get_gpu(self, gpu_id):