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

2020-04-21 14:29:15 -04:00 · 2020-04-21 14:29:15 -04:00 · bafdeca42f
parent 29c7d2f195
commit bafdeca42f
1 changed files with 1 additions and 1 deletions
--- a/pytorch_lightning/trainer/distrib_data_parallel.py
+++ b/pytorch_lightning/trainer/distrib_data_parallel.py
@ -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)