Add a reference to the Trainer on the LightningDataModule

pytorch_lightning/trainer/connectors/data_connector.py

@@ -14,7 +14,7 @@
 
 from pytorch_lightning.core.datamodule import LightningDataModule
 from pytorch_lightning.utilities.exceptions import MisconfigurationException
-from typing import List, Union
+from typing import List, Optional, Union
 from torch.utils.data import DataLoader
 from pytorch_lightning.utilities.model_utils import is_overridden
 
@@ -101,7 +101,7 @@ def attach_dataloaders(self, model, train_dataloader=None, val_dataloaders=None,
         if test_dataloaders is not None:
             model.test_dataloader = _PatchDataLoader(test_dataloaders)
 
-    def attach_datamodule(self, model, datamodule, stage):
+    def attach_datamodule(self, model, datamodule: Optional[LightningDataModule], stage: str) -> None:
 
         # We use datamodule if it's been provided on .fit or .test, otherwise we check model for it
         datamodule = datamodule or getattr(model, 'datamodule', None)
@@ -122,6 +122,7 @@ def attach_datamodule(self, model, datamodule, stage):
                 model.transfer_batch_to_device = datamodule.transfer_batch_to_device
 
             self.trainer.datamodule = datamodule
+            datamodule.trainer = self.trainer
 
 
 class _PatchDataLoader(object):

tests/core/test_datamodules.py

@@ -255,6 +255,32 @@ def test_full_loop(tmpdir):
     assert result['test_acc'] > 0.8
 
 
+def test_trainer_attached_to_dm(tmpdir):
+    reset_seed()
+
+    dm = TrialMNISTDataModule(tmpdir)
+
+    model = EvalModelTemplate()
+
+    trainer = Trainer(
+        default_root_dir=tmpdir,
+        max_epochs=3,
+        weights_summary=None,
+        deterministic=True,
+    )
+
+    # fit model
+    result = trainer.fit(model, dm)
+    assert result == 1
+    assert dm.trainer is not None
+
+    # test
+    result = trainer.test(datamodule=dm)
+    result = result[0]
+    assert dm.trainer is not None
+
+
+
 @pytest.mark.skipif(torch.cuda.device_count() < 1, reason="test requires multi-GPU machine")
 def test_full_loop_single_gpu(tmpdir):
     reset_seed()