From 7b6c6c5fc1adf4fc685125deccb567e4c1bd4f4e Mon Sep 17 00:00:00 2001 From: adamoudad Date: Mon, 30 Nov 2020 08:13:21 +0900 Subject: [PATCH] Fix training for single GPU --- train.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/train.py b/train.py index 6d7b618..1c97d37 100644 --- 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) )