Fix training for single GPU

jason9693 · Nov 29, 2020 · 7b6c6c5 · 7b6c6c5
1 parent 5f18337
commit 7b6c6c5
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/train.py b/train.py
@@ -135,7 +135,8 @@
 
         # switch output device to: gpu-1 ~ gpu-n
         sw_start = time.time()
-        mt.output_device = idx % (torch.cuda.device_count() -1) + 1
+        if torch.cuda.device_count() > 1:
+            mt.output_device = idx % (torch.cuda.device_count() -1) + 1
         sw_end = time.time()
         if config.debug:
             print('output switch time: {}'.format(sw_end - sw_start) )