more robust handling of gpu parameters
@@ -26,6 +26,8 @@ def torch_setup(gpus=None, random_seed=None):
if torch.cuda.is_available():
if gpus is None:
gpus = range(torch.cuda.device_count()) if torch.cuda.is_available() else None
+ elif len(gpus) == 0 or gpus[0] == -1:
+ gpus=None
else:
max_device_index = torch.cuda.device_count() - 1
@@ -38,12 +40,12 @@ def torch_setup(gpus=None, random_seed=None):
if len(gpus) == 0:
gpus = None
else:
- if len(gpus) > 0:
+ if gpus is not None and len(gpus) > 0 and gpus[0] != -1:
raise CudaNotAvailableError()
gpus = None
- if len(gpus) > 1:
+ if gpus is not None and len(gpus) > 1:
raise MultiGpuNotSupportedError()
if random_seed is not None:
