소스 검색

Allow using a single GPU

kaiix 9 년 전
부모
커밋
dab1517e74
1개의 변경된 파일2개의 추가작업 그리고 1개의 파일을 삭제
  1. 2 1
      textsum/seq2seq_attention_model.py

+ 2 - 1
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):