🐛 .

Lightning-AI · Aug 1, 2020 · f059cc4 · f059cc4
1 parent 1b77442
commit f059cc4
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 3 deletions.
diff --git a/pytorch_lightning/trainer/distrib_data_parallel.py b/pytorch_lightning/trainer/distrib_data_parallel.py
@@ -488,7 +488,7 @@ def spawn_ddp_children(self, model):
 
         return results
 
-    def ddp_train(self, process_idx, mp_queue, model, is_master=False, proc_offset=0, datamodule=None):
+    def ddp_train(self, process_idx, mp_queue, model, is_master=False, proc_offset=0):
         """
         Entry point for ddp
 
@@ -531,8 +531,8 @@ def ddp_train(self, process_idx, mp_queue, model, is_master=False, proc_offset=0
 
         # call setup after the ddp process has connected
         if not self.testing:
-            if datamodule is not None:
-                datamodule.setup('fit')
+            if self.datamodule is not None:
+                self.datamodule.setup('fit')
             self.setup('fit')
             model.setup('fit')
 

diff --git a/pytorch_lightning/trainer/trainer.py b/pytorch_lightning/trainer/trainer.py
@@ -378,6 +378,7 @@ def __init__(
 
         # training state
         self.model = None
+        self.datamodule = None
         self.testing = False
         self.prepare_data_per_node = prepare_data_per_node
         self.lr_schedulers = []