🚧 wip

Lightning-AI · Jul 31, 2020 · 8562b3a · 8562b3a
1 parent 1ed4024
commit 8562b3a
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/pytorch_lightning/trainer/distrib_data_parallel.py b/pytorch_lightning/trainer/distrib_data_parallel.py
@@ -143,6 +143,7 @@ def train_fx(trial_hparams, cluster_manager, _):
 from pytorch_lightning.utilities import NATIVE_AMP_AVALAIBLE
 from pytorch_lightning.utilities.exceptions import MisconfigurationException
 from pytorch_lightning.utilities.distributed import rank_zero_only, rank_zero_warn, rank_zero_info
+from pytorch_lightning.core.datamodule import LightningDataModule
 from pytorch_lightning.core.lightning import LightningModule
 
 
@@ -202,6 +203,7 @@ class TrainerDDPMixin(ABC):
     node_rank: int
     tpu_cores: int
     testing: bool
+    datamodule: Optional[LightningDataModule]
 
     @property
     @abstractmethod