Skip to content

Commit

Permalink
Multi-GPU default to single device 0
Browse files Browse the repository at this point in the history
  • Loading branch information
glenn-jocher committed Jun 9, 2021
1 parent 958ab92 commit cfd4324
Showing 1 changed file with 1 addition and 1 deletion.
2 changes: 1 addition & 1 deletion utils/torch_utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -72,7 +72,7 @@ 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}'
Expand Down

0 comments on commit cfd4324

Please sign in to comment.