mirrors
/
yolov5
mirror of https://github.com/ultralytics/yolov5.git
1
0
Fork 0
Code Issues Projects Releases Wiki Activity

Update torch_utils.py (#1895)

pull/1897/head
Glenn Jocher 2021-01-10 12:26:29 -08:00 committed by GitHub
parent 0f11aaf551
commit ffef77124e
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 1 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
utils/torch_utils.py
View File

@ -61,7 +61,7 @@ def select_device(device='', batch_size=None):
os.environ['CUDA_VISIBLE_DEVICES'] = device # set environment variable
assert torch.cuda.is_available(), f'CUDA unavailable, invalid device {device} requested' # check availability
cuda = torch.cuda.is_available() and not cpu
cuda = not cpu and torch.cuda.is_available()
if cuda:
n = torch.cuda.device_count()
if n > 1 and batch_size: # check that batch_size is compatible with device_count