wassname/pytorch-lightning
1
0
Fork 0
mirror of https://github.com/wassname/pytorch-lightning.git synced 2026-09-10 12:21:57 +08:00
Code Issues Packages Projects Releases Wiki Activity

moved cuda flags inside trainer

Browse Source
This commit is contained in:
William Falcon
2019-07-08 10:00:04 -04:00
parent 79ca5f6265
commit cdbbf9abe3
1 changed files with 5 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files
pytorch_lightning/models/trainer.py
+5 -3
View File
@@ -91,16 +91,18 @@ 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()))
else:
self.data_parallel_device_ids = [int(x.strip()) for x in gpus.split(',')]
# set the correct cuda visible devices (using pci order)
os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
os.environ["CUDA_VISIBLE_DEVICES"] = ','.join([str(x) for x in self.data_parallel_device_ids])
self.data_parallel = self.data_parallel_device_ids is not None and len(self.data_parallel_device_ids) > 0
# set the correct cuda visible devices (using pci order)
os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
os.environ["CUDA_VISIBLE_DEVICES"] = ','.join(self.data_parallel_device_ids)
# process info
self.proc_rank = 0