Update torch_utils.py

ultralytics · Jan 26, 2022 · 24b80ca · 24b80ca
1 parent 7b98fa6
commit 24b80ca
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/utils/torch_utils.py b/utils/torch_utils.py
@@ -62,9 +62,9 @@ def select_device(device='', batch_size=0, newline=True):
         os.environ['CUDA_VISIBLE_DEVICES'] = '-1'  # force torch.cuda.is_available() = False
     elif device:  # non-cpu device requested
         nd = len(os.getenv('CUDA_VISIBLE_DEVICES', '').replace(',',''))  # number of CUDA devices
-        assert torch.cuda.is_available(), 'CUDA is not available, use `--device cpu` or do not pass a --device'
         assert nd > int(max(device.split(','))), f'Invalid `--device {device}` request, valid devices are 0 - {nd - 1}'
-        os.environ['CUDA_VISIBLE_DEVICES'] = device  # set environment variable (must be after asserts)
+        os.environ['CUDA_VISIBLE_DEVICES'] = device  # set environment variable (before assert cuda.is_available)
+        assert torch.cuda.is_available(), 'CUDA is not available, use `--device cpu` or do not pass a --device'
 
     cuda = not cpu and torch.cuda.is_available()
     if cuda: