From e4baa353e2e6d12b4c041e9cd09e5598ee445418 Mon Sep 17 00:00:00 2001 From: Glenn Jocher Date: Sun, 10 Jan 2021 11:59:12 -0800 Subject: [PATCH] Update torch_utils.py --- utils/torch_utils.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/utils/torch_utils.py b/utils/torch_utils.py index e441b8e0626a..2cb09e71ce71 100644 --- a/utils/torch_utils.py +++ b/utils/torch_utils.py @@ -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