From 235781564011cdb880c66dafe87cbada7edb709b Mon Sep 17 00:00:00 2001 From: William Falcon Date: Sun, 21 Jul 2019 08:08:21 -0400 Subject: [PATCH] release v0.3 --- pytorch_lightning/models/trainer.py | 11 ++++++++--- setup.py | 2 +- 2 files changed, 9 insertions(+), 4 deletions(-) diff --git a/pytorch_lightning/models/trainer.py b/pytorch_lightning/models/trainer.py index 9cf75f4c..ba47ab6c 100644 --- a/pytorch_lightning/models/trainer.py +++ b/pytorch_lightning/models/trainer.py @@ -154,10 +154,15 @@ class Trainer(TrainerIO): # if gpus = -1 then use all available devices # otherwise, split the string using commas if gpus is not None: - if gpus == '-1': - self.data_parallel_device_ids = list(range(0, torch.cuda.device_count())) + if type(gpus) is list: + self.data_parallel_device_ids = gpus + elif type(gpus) is str: + if gpus == '-1': + self.data_parallel_device_ids = list(range(0, torch.cuda.device_count())) + else: + self.data_parallel_device_ids = [int(x.strip()) for x in gpus.split(',')] else: - self.data_parallel_device_ids = [int(x.strip()) for x in gpus.split(',')] + raise Exception('gpus has to be a string or list of ids') # set the correct cuda visible devices (using pci order) os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID" diff --git a/setup.py b/setup.py index 59f1866f..f23eaa4b 100755 --- a/setup.py +++ b/setup.py @@ -7,7 +7,7 @@ from setuptools import setup, find_packages # http://blog.ionelmc.ro/2014/05/25/python-packaging/ setup( name="pytorch-lightning", - version='0.2.6', + version='0.3', description="The Keras for ML researchers using PyTorch", author="William Falcon", author_email="waf2107@columbia.edu",