Spaces:

xiang-wuu
/

yolov5

Runtime error

glenn-jocher commited on Jan 10, 2021

Commit

ffef771

•

1 Parent(s): 0f11aaf

Update torch_utils.py (#1895)

Files changed (1) hide show

utils/torch_utils.py CHANGED Viewed

@@ -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
     if cuda:
         n = torch.cuda.device_count()
         if n > 1 and batch_size:  # check that batch_size is compatible with device_count

         os.environ['CUDA_VISIBLE_DEVICES'] = device  # set environment variable
         assert torch.cuda.is_available(), f'CUDA unavailable, invalid device {device} requested'  # check availability
+    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