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

Replace GPU device idx with current process index (#1541)

Browse Source
This commit is contained in:
Kevin Chen
2020-04-21 14:29:15 -04:00
committed by GitHub
parent 29c7d2f195
commit bafdeca42f
1 changed files with 1 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files
pytorch_lightning/trainer/distrib_data_parallel.py
+1 -1
View File
@@ -327,7 +327,7 @@ class TrainerDDPMixin(ABC):
# MODEL
# copy model to each gpu
if self.on_gpu:
self.root_gpu = self.data_parallel_device_ids[process_idx]
self.root_gpu = process_idx
torch.cuda.set_device(self.root_gpu)
model.cuda(self.root_gpu)