Skip to content

GitLab

Commit dab1517e authored by kaiix's avatar kaiix
Browse files

Allow using a single GPU

parent 761c59c4
Hide whitespace changes
Inline Side-by-side
Showing with 2 additions and 1 deletion
textsum/seq2seq_attention_model.py
View file @ dab1517e
...@@ -105,7 +105,8 @@ class Seq2SeqAttentionModel(object): ...@@ -105,7 +105,8 @@ class Seq2SeqAttentionModel(object):
if self._num_gpus == 0: if self._num_gpus == 0:
return '' return ''
dev = '/gpu:%d' % self._cur_gpu 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 return dev
def _get_gpu(self, gpu_id): def _get_gpu(self, gpu_id):
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment