apoorva/greenhouse
2
1
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

updates

Browse Source
This commit is contained in:
Glenn Jocher
2019-03-04 15:46:25 +01:00
parent 66211d1147
commit 65ed0d1122
2 changed files with 5 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files
utils/torch_utils.py
+3 -2
View File
@@ -15,8 +15,9 @@ def select_device(force_cpu=False):
cuda = torch.cuda.is_available()
device = torch.device('cuda:0' if cuda else 'cpu')
# if torch.cuda.device_count() > 1:
# print('WARNING Using GPU0 Only: https://github.com/ultralytics/yolov3/issues/21')
if torch.cuda.device_count() > 1:
print('Found %g GPUs' % torch.cuda.device_count())
print('WARNING Using GPU0 Only: https://github.com/ultralytics/yolov3/issues/21')
# torch.cuda.set_device(0) # OPTIONAL: Set your GPU if multiple available
# # print('Using ', torch.cuda.device_count(), ' GPUs')