From 3f4de4518ba4e7a94fa79b3cfa47db062bd6c2d1 Mon Sep 17 00:00:00 2001 From: Glenn Jocher Date: Wed, 9 Jun 2021 15:31:47 +0200 Subject: [PATCH] Multi-GPU default to single device 0 (#3554) * Multi-GPU default to single device 0 * Multi-GPU default to single device 0 * add space (cherry picked from commit 8b5086c21ba227c0257d94ea34cb46124a9c559a) --- utils/torch_utils.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/utils/torch_utils.py b/utils/torch_utils.py index 6a7d07634813..b690dbe96700 100644 --- a/utils/torch_utils.py +++ b/utils/torch_utils.py @@ -72,11 +72,11 @@ def select_device(device='', batch_size=None): cuda = not cpu and torch.cuda.is_available() if cuda: - devices = device.split(',') if device else range(torch.cuda.device_count()) # i.e. 0,1,6,7 + devices = device.split(',') if device else '0' # range(torch.cuda.device_count()) # i.e. 0,1,6,7 n = len(devices) # device count if n > 1 and batch_size: # check batch_size is divisible by device_count assert batch_size % n == 0, f'batch-size {batch_size} not multiple of GPU count {n}' - space = ' ' * len(s) + space = ' ' * (len(s) + 1) for i, d in enumerate(devices): p = torch.cuda.get_device_properties(i) s += f"{'' if i == 0 else space}CUDA:{d} ({p.name}, {p.total_memory / 1024 ** 2}MB)\n" # bytes to MB