Ankle/yolov5_birdproject
1
0 Fork 0
Code Issues Pull Requests Actions 3 Packages Projects Releases Wiki Activity

Update torch_utils.py (#1895)

Browse Source
This commit is contained in:
Glenn Jocher authored and GitHub committed 2021-01-10 12:26:29 -08:00
1 parent 0f11aaf551
commit ffef77124e
1 file changed
+1 -1
Download Patch File Download Diff File Expand all files Collapse all files
utils/torch_utils.py
+1 -1
View File
@@ -61,7 +61,7 @@ def select_device(device='', batch_size=None):
os.environ['CUDA_VISIBLE_DEVICES'] = device # set environment variable os.environ['CUDA_VISIBLE_DEVICES'] = device # set environment variable
assert torch.cuda.is_available(), f'CUDA unavailable, invalid device {device} requested' # check availability 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: if cuda:
n = torch.cuda.device_count() n = torch.cuda.device_count()
if n > 1 and batch_size: # check that batch_size is compatible with device_count if n > 1 and batch_size: # check that batch_size is compatible with device_count