We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
DataModules should have a reference to the Trainer
## 🚀 Feature
The DataModule should have a reference to the Trainer, so attributes on the Trainer can be accessed during prepare_data, setup, and the dataloader methods. 

cc @nateraw 

### Motivation
Further uncouple datamodules and lightning modules. For instance, it can be useful for DataModules to know what the current epoch is during training. If I want to use the DataModules, I don't have access to this. However this is readily available if I define my `train_dataloader` in the LightningModule instead. 


### Pitch
We can follow the same setup that LightningModules have and how the trainer is initialized for them


--- END ISSUE ---

Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
--- BEGIN FILES ---
Path: `pytorch_lightning/trainer/connectors/data_connector.py`
Content:
```
1 # Copyright The PyTorch Lightning team.
2 #
3 # Licensed under the Apache License, Version 2.0 (the "License");
4 # you may not use this file except in compliance with the License.
5 # You may obtain a copy of the License at
6 #
7 #     http://www.apache.org/licenses/LICENSE-2.0
8 #
9 # Unless required by applicable law or agreed to in writing, software
10 # distributed under the License is distributed on an "AS IS" BASIS,
11 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12 # See the License for the specific language governing permissions and
13 # limitations under the License.
14 
15 from pytorch_lightning.core.datamodule import LightningDataModule
16 from pytorch_lightning.utilities.exceptions import MisconfigurationException
17 from typing import List, Union
18 from torch.utils.data import DataLoader
19 from pytorch_lightning.utilities.model_utils import is_overridden
20 
21 
22 class DataConnector(object):
23 
24     def __init__(self, trainer):
25         self.trainer = trainer
26 
27     def on_trainer_init(self, check_val_every_n_epoch, reload_dataloaders_every_epoch, prepare_data_per_node):
28         self.trainer.datamodule = None
29         self.trainer.prepare_data_per_node = prepare_data_per_node
30 
31         self.trainer.check_val_every_n_epoch = check_val_every_n_epoch
32         self.trainer.reload_dataloaders_every_epoch = reload_dataloaders_every_epoch
33         self.trainer._is_data_prepared = False
34 
35     def get_profiled_train_dataloader(self, train_dataloader):
36         profiled_dl = self.trainer.profiler.profile_iterable(
37             enumerate(self._with_is_last(train_dataloader)),
38             "get_train_batch"
39         )
40         return profiled_dl
41 
42     def _with_is_last(self, iterable):
43         """Pass through values from the given iterable with an added boolean indicating if this is the last item.
44         See `https://stackoverflow.com/a/1630350 <https://stackoverflow.com/a/1630350>`_"""
45         it = iter(iterable)
46         last = next(it)
47         for val in it:
48             # yield last and has next
49             yield last, False
50             last = val
51         # yield last, no longer has next
52         yield last, True
53 
54     def prepare_data(self, model):
55         # on multi-gpu jobs we only want to manipulate (download, etc) on node_rank=0, local_rank=0
56         # or in the case where each node needs to do its own manipulation in which case just local_rank=0
57         if self.can_prepare_data():
58             if self.trainer.datamodule is not None:
59                 self.trainer.datamodule.prepare_data()
60             model.prepare_data()
61             self.trainer._is_data_prepared = True
62 
63     def can_prepare_data(self):
64         should_call_dm_prepare_data = True
65         if self.trainer.datamodule is not None and is_overridden('prepare_data', self.trainer.datamodule):
66             should_call_dm_prepare_data = not self.trainer.datamodule.has_prepared_data
67 
68         if self.trainer.prepare_data_per_node:
69             return self.trainer.local_rank == 0 and should_call_dm_prepare_data
70         else:
71             return self.trainer.node_rank == 0 and self.trainer.local_rank == 0 and should_call_dm_prepare_data
72 
73     def attach_data(self, model, train_dataloader, val_dataloaders, datamodule):
74         # if a datamodule comes in as the second arg, then fix it for the user
75         if isinstance(train_dataloader, LightningDataModule):
76             datamodule = train_dataloader
77             train_dataloader = None
78 
79         self.__enforce_datamodule_dataloader_override(train_dataloader, val_dataloaders, datamodule)
80 
81         # set up the passed in dataloaders (if needed)
82         self.attach_dataloaders(model, train_dataloader, val_dataloaders)
83         self.attach_datamodule(model, datamodule, 'fit')
84 
85     def __enforce_datamodule_dataloader_override(self, train_dataloader, val_dataloaders, datamodule):
86         # If you supply a datamodule you can't supply train_dataloader or val_dataloaders
87         if (train_dataloader is not None or val_dataloaders is not None) and datamodule is not None:
88             raise MisconfigurationException(
89                 'You cannot pass train_dataloader or val_dataloaders to trainer.fit if you supply a datamodule'
90             )
91 
92     def attach_dataloaders(self, model, train_dataloader=None, val_dataloaders=None, test_dataloaders=None):
93         # when dataloader is passed via fit, patch the train_dataloader
94         # functions to overwrite with these implementations
95         if train_dataloader is not None:
96             model.train_dataloader = _PatchDataLoader(train_dataloader)
97 
98         if val_dataloaders is not None:
99             model.val_dataloader = _PatchDataLoader(val_dataloaders)
100 
101         if test_dataloaders is not None:
102             model.test_dataloader = _PatchDataLoader(test_dataloaders)
103 
104     def attach_datamodule(self, model, datamodule, stage):
105 
106         # We use datamodule if it's been provided on .fit or .test, otherwise we check model for it
107         datamodule = datamodule or getattr(model, 'datamodule', None)
108 
109         # If we have a datamodule, attach necessary hooks + dataloaders
110         if datamodule:
111 
112             # Override loader hooks
113             if is_overridden('train_dataloader', datamodule):
114                 model.train_dataloader = datamodule.train_dataloader
115             if is_overridden('val_dataloader', datamodule):
116                 model.val_dataloader = datamodule.val_dataloader
117             if is_overridden('test_dataloader', datamodule):
118                 model.test_dataloader = datamodule.test_dataloader
119 
120             # Override transfer_batch_to_device if dataset-specific to_device logic has been defined in datamodule
121             if is_overridden('transfer_batch_to_device', datamodule):
122                 model.transfer_batch_to_device = datamodule.transfer_batch_to_device
123 
124             self.trainer.datamodule = datamodule
125 
126 
127 class _PatchDataLoader(object):
128     r"""
129     Callable object for patching dataloaders passed into trainer.fit().
130     Use this class to override model.*_dataloader() and be pickle-compatible.
131 
132     Args:
133         dataloader: Dataloader object to return when called.
134 
135     """
136 
137     def __init__(self, dataloader: Union[List[DataLoader], DataLoader]):
138         self.dataloader = dataloader
139 
140         # cannot pickle __code__ so cannot verify if PatchDataloader
141         # exists which shows dataloader methods have been overwritten.
142         # so, we hack it by using the string representation
143         self.patch_loader_code = str(self.__call__.__code__)
144 
145     def __call__(self) -> Union[List[DataLoader], DataLoader]:
146         return self.dataloader
147
```
--- END FILES ---

Please first localize the bug based on the issue statement, and then generate a patch according to the `git diff` format fenced by three backticks.

Here is an example:

```diff
diff --git a/examples/server_async.py b/examples/server_async.py
--- a/examples/server_async.py
+++ b/examples/server_async.py
@@ -313,4 +313,4 @@
 
 
 if __name__ == "__main__":
-    asyncio.run(run_async_server("."), debug=True)
+    asyncio.run(run_async_server(), debug=True)
diff --git a/examples/server_sync.py b/examples/server_sync.py
--- a/examples/server_sync.py
+++ b/examples/server_sync.py
@@ -313,5 +313,5 @@
 
 
 if __name__ == "__main__":
-    server = run_sync_server(".")
+    server = run_sync_server()
     server.shutdown()

```

diff --git a/pytorch_lightning/trainer/connectors/data_connector.py b/pytorch_lightning/trainer/connectors/data_connector.py
--- a/pytorch_lightning/trainer/connectors/data_connector.py
+++ b/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 @@
         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 @@
                 model.transfer_batch_to_device = datamodule.transfer_batch_to_device
 
             self.trainer.datamodule = datamodule
+            datamodule.trainer = self.trainer
 
 
 class _PatchDataLoader(object):

{"golden_diff": "diff --git a/pytorch_lightning/trainer/connectors/data_connector.py b/pytorch_lightning/trainer/connectors/data_connector.py\n--- a/pytorch_lightning/trainer/connectors/data_connector.py\n+++ b/pytorch_lightning/trainer/connectors/data_connector.py\n@@ -14,7 +14,7 @@\n \n from pytorch_lightning.core.datamodule import LightningDataModule\n from pytorch_lightning.utilities.exceptions import MisconfigurationException\n-from typing import List, Union\n+from typing import List, Optional, Union\n from torch.utils.data import DataLoader\n from pytorch_lightning.utilities.model_utils import is_overridden\n \n@@ -101,7 +101,7 @@\n         if test_dataloaders is not None:\n             model.test_dataloader = _PatchDataLoader(test_dataloaders)\n \n-    def attach_datamodule(self, model, datamodule, stage):\n+    def attach_datamodule(self, model, datamodule: Optional[LightningDataModule], stage: str) -> None:\n \n         # We use datamodule if it's been provided on .fit or .test, otherwise we check model for it\n         datamodule = datamodule or getattr(model, 'datamodule', None)\n@@ -122,6 +122,7 @@\n                 model.transfer_batch_to_device = datamodule.transfer_batch_to_device\n \n             self.trainer.datamodule = datamodule\n+            datamodule.trainer = self.trainer\n \n \n class _PatchDataLoader(object):\n", "issue": "DataModules should have a reference to the Trainer\n## \ud83d\ude80 Feature\r\nThe DataModule should have a reference to the Trainer, so attributes on the Trainer can be accessed during prepare_data, setup, and the dataloader methods. \r\n\r\ncc @nateraw \r\n\r\n### Motivation\r\nFurther uncouple datamodules and lightning modules. For instance, it can be useful for DataModules to know what the current epoch is during training. If I want to use the DataModules, I don't have access to this. However this is readily available if I define my `train_dataloader` in the LightningModule instead. \r\n\r\n\r\n### Pitch\r\nWe can follow the same setup that LightningModules have and how the trainer is initialized for them\r\n\n", "before_files": [{"content": "# Copyright The PyTorch Lightning team.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nfrom pytorch_lightning.core.datamodule import LightningDataModule\nfrom pytorch_lightning.utilities.exceptions import MisconfigurationException\nfrom typing import List, Union\nfrom torch.utils.data import DataLoader\nfrom pytorch_lightning.utilities.model_utils import is_overridden\n\n\nclass DataConnector(object):\n\n    def __init__(self, trainer):\n        self.trainer = trainer\n\n    def on_trainer_init(self, check_val_every_n_epoch, reload_dataloaders_every_epoch, prepare_data_per_node):\n        self.trainer.datamodule = None\n        self.trainer.prepare_data_per_node = prepare_data_per_node\n\n        self.trainer.check_val_every_n_epoch = check_val_every_n_epoch\n        self.trainer.reload_dataloaders_every_epoch = reload_dataloaders_every_epoch\n        self.trainer._is_data_prepared = False\n\n    def get_profiled_train_dataloader(self, train_dataloader):\n        profiled_dl = self.trainer.profiler.profile_iterable(\n            enumerate(self._with_is_last(train_dataloader)),\n            \"get_train_batch\"\n        )\n        return profiled_dl\n\n    def _with_is_last(self, iterable):\n        \"\"\"Pass through values from the given iterable with an added boolean indicating if this is the last item.\n        See `https://stackoverflow.com/a/1630350 <https://stackoverflow.com/a/1630350>`_\"\"\"\n        it = iter(iterable)\n        last = next(it)\n        for val in it:\n            # yield last and has next\n            yield last, False\n            last = val\n        # yield last, no longer has next\n        yield last, True\n\n    def prepare_data(self, model):\n        # on multi-gpu jobs we only want to manipulate (download, etc) on node_rank=0, local_rank=0\n        # or in the case where each node needs to do its own manipulation in which case just local_rank=0\n        if self.can_prepare_data():\n            if self.trainer.datamodule is not None:\n                self.trainer.datamodule.prepare_data()\n            model.prepare_data()\n            self.trainer._is_data_prepared = True\n\n    def can_prepare_data(self):\n        should_call_dm_prepare_data = True\n        if self.trainer.datamodule is not None and is_overridden('prepare_data', self.trainer.datamodule):\n            should_call_dm_prepare_data = not self.trainer.datamodule.has_prepared_data\n\n        if self.trainer.prepare_data_per_node:\n            return self.trainer.local_rank == 0 and should_call_dm_prepare_data\n        else:\n            return self.trainer.node_rank == 0 and self.trainer.local_rank == 0 and should_call_dm_prepare_data\n\n    def attach_data(self, model, train_dataloader, val_dataloaders, datamodule):\n        # if a datamodule comes in as the second arg, then fix it for the user\n        if isinstance(train_dataloader, LightningDataModule):\n            datamodule = train_dataloader\n            train_dataloader = None\n\n        self.__enforce_datamodule_dataloader_override(train_dataloader, val_dataloaders, datamodule)\n\n        # set up the passed in dataloaders (if needed)\n        self.attach_dataloaders(model, train_dataloader, val_dataloaders)\n        self.attach_datamodule(model, datamodule, 'fit')\n\n    def __enforce_datamodule_dataloader_override(self, train_dataloader, val_dataloaders, datamodule):\n        # If you supply a datamodule you can't supply train_dataloader or val_dataloaders\n        if (train_dataloader is not None or val_dataloaders is not None) and datamodule is not None:\n            raise MisconfigurationException(\n                'You cannot pass train_dataloader or val_dataloaders to trainer.fit if you supply a datamodule'\n            )\n\n    def attach_dataloaders(self, model, train_dataloader=None, val_dataloaders=None, test_dataloaders=None):\n        # when dataloader is passed via fit, patch the train_dataloader\n        # functions to overwrite with these implementations\n        if train_dataloader is not None:\n            model.train_dataloader = _PatchDataLoader(train_dataloader)\n\n        if val_dataloaders is not None:\n            model.val_dataloader = _PatchDataLoader(val_dataloaders)\n\n        if test_dataloaders is not None:\n            model.test_dataloader = _PatchDataLoader(test_dataloaders)\n\n    def attach_datamodule(self, model, datamodule, stage):\n\n        # We use datamodule if it's been provided on .fit or .test, otherwise we check model for it\n        datamodule = datamodule or getattr(model, 'datamodule', None)\n\n        # If we have a datamodule, attach necessary hooks + dataloaders\n        if datamodule:\n\n            # Override loader hooks\n            if is_overridden('train_dataloader', datamodule):\n                model.train_dataloader = datamodule.train_dataloader\n            if is_overridden('val_dataloader', datamodule):\n                model.val_dataloader = datamodule.val_dataloader\n            if is_overridden('test_dataloader', datamodule):\n                model.test_dataloader = datamodule.test_dataloader\n\n            # Override transfer_batch_to_device if dataset-specific to_device logic has been defined in datamodule\n            if is_overridden('transfer_batch_to_device', datamodule):\n                model.transfer_batch_to_device = datamodule.transfer_batch_to_device\n\n            self.trainer.datamodule = datamodule\n\n\nclass _PatchDataLoader(object):\n    r\"\"\"\n    Callable object for patching dataloaders passed into trainer.fit().\n    Use this class to override model.*_dataloader() and be pickle-compatible.\n\n    Args:\n        dataloader: Dataloader object to return when called.\n\n    \"\"\"\n\n    def __init__(self, dataloader: Union[List[DataLoader], DataLoader]):\n        self.dataloader = dataloader\n\n        # cannot pickle __code__ so cannot verify if PatchDataloader\n        # exists which shows dataloader methods have been overwritten.\n        # so, we hack it by using the string representation\n        self.patch_loader_code = str(self.__call__.__code__)\n\n    def __call__(self) -> Union[List[DataLoader], DataLoader]:\n        return self.dataloader\n", "path": "pytorch_lightning/trainer/connectors/data_connector.py"}], "after_files": [{"content": "# Copyright The PyTorch Lightning team.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nfrom pytorch_lightning.core.datamodule import LightningDataModule\nfrom pytorch_lightning.utilities.exceptions import MisconfigurationException\nfrom typing import List, Optional, Union\nfrom torch.utils.data import DataLoader\nfrom pytorch_lightning.utilities.model_utils import is_overridden\n\n\nclass DataConnector(object):\n\n    def __init__(self, trainer):\n        self.trainer = trainer\n\n    def on_trainer_init(self, check_val_every_n_epoch, reload_dataloaders_every_epoch, prepare_data_per_node):\n        self.trainer.datamodule = None\n        self.trainer.prepare_data_per_node = prepare_data_per_node\n\n        self.trainer.check_val_every_n_epoch = check_val_every_n_epoch\n        self.trainer.reload_dataloaders_every_epoch = reload_dataloaders_every_epoch\n        self.trainer._is_data_prepared = False\n\n    def get_profiled_train_dataloader(self, train_dataloader):\n        profiled_dl = self.trainer.profiler.profile_iterable(\n            enumerate(self._with_is_last(train_dataloader)),\n            \"get_train_batch\"\n        )\n        return profiled_dl\n\n    def _with_is_last(self, iterable):\n        \"\"\"Pass through values from the given iterable with an added boolean indicating if this is the last item.\n        See `https://stackoverflow.com/a/1630350 <https://stackoverflow.com/a/1630350>`_\"\"\"\n        it = iter(iterable)\n        last = next(it)\n        for val in it:\n            # yield last and has next\n            yield last, False\n            last = val\n        # yield last, no longer has next\n        yield last, True\n\n    def prepare_data(self, model):\n        # on multi-gpu jobs we only want to manipulate (download, etc) on node_rank=0, local_rank=0\n        # or in the case where each node needs to do its own manipulation in which case just local_rank=0\n        if self.can_prepare_data():\n            if self.trainer.datamodule is not None:\n                self.trainer.datamodule.prepare_data()\n            model.prepare_data()\n            self.trainer._is_data_prepared = True\n\n    def can_prepare_data(self):\n        should_call_dm_prepare_data = True\n        if self.trainer.datamodule is not None and is_overridden('prepare_data', self.trainer.datamodule):\n            should_call_dm_prepare_data = not self.trainer.datamodule.has_prepared_data\n\n        if self.trainer.prepare_data_per_node:\n            return self.trainer.local_rank == 0 and should_call_dm_prepare_data\n        else:\n            return self.trainer.node_rank == 0 and self.trainer.local_rank == 0 and should_call_dm_prepare_data\n\n    def attach_data(self, model, train_dataloader, val_dataloaders, datamodule):\n        # if a datamodule comes in as the second arg, then fix it for the user\n        if isinstance(train_dataloader, LightningDataModule):\n            datamodule = train_dataloader\n            train_dataloader = None\n\n        self.__enforce_datamodule_dataloader_override(train_dataloader, val_dataloaders, datamodule)\n\n        # set up the passed in dataloaders (if needed)\n        self.attach_dataloaders(model, train_dataloader, val_dataloaders)\n        self.attach_datamodule(model, datamodule, 'fit')\n\n    def __enforce_datamodule_dataloader_override(self, train_dataloader, val_dataloaders, datamodule):\n        # If you supply a datamodule you can't supply train_dataloader or val_dataloaders\n        if (train_dataloader is not None or val_dataloaders is not None) and datamodule is not None:\n            raise MisconfigurationException(\n                'You cannot pass train_dataloader or val_dataloaders to trainer.fit if you supply a datamodule'\n            )\n\n    def attach_dataloaders(self, model, train_dataloader=None, val_dataloaders=None, test_dataloaders=None):\n        # when dataloader is passed via fit, patch the train_dataloader\n        # functions to overwrite with these implementations\n        if train_dataloader is not None:\n            model.train_dataloader = _PatchDataLoader(train_dataloader)\n\n        if val_dataloaders is not None:\n            model.val_dataloader = _PatchDataLoader(val_dataloaders)\n\n        if test_dataloaders is not None:\n            model.test_dataloader = _PatchDataLoader(test_dataloaders)\n\n    def attach_datamodule(self, model, datamodule: Optional[LightningDataModule], stage: str) -> None:\n\n        # We use datamodule if it's been provided on .fit or .test, otherwise we check model for it\n        datamodule = datamodule or getattr(model, 'datamodule', None)\n\n        # If we have a datamodule, attach necessary hooks + dataloaders\n        if datamodule:\n\n            # Override loader hooks\n            if is_overridden('train_dataloader', datamodule):\n                model.train_dataloader = datamodule.train_dataloader\n            if is_overridden('val_dataloader', datamodule):\n                model.val_dataloader = datamodule.val_dataloader\n            if is_overridden('test_dataloader', datamodule):\n                model.test_dataloader = datamodule.test_dataloader\n\n            # Override transfer_batch_to_device if dataset-specific to_device logic has been defined in datamodule\n            if is_overridden('transfer_batch_to_device', datamodule):\n                model.transfer_batch_to_device = datamodule.transfer_batch_to_device\n\n            self.trainer.datamodule = datamodule\n            datamodule.trainer = self.trainer\n\n\nclass _PatchDataLoader(object):\n    r\"\"\"\n    Callable object for patching dataloaders passed into trainer.fit().\n    Use this class to override model.*_dataloader() and be pickle-compatible.\n\n    Args:\n        dataloader: Dataloader object to return when called.\n\n    \"\"\"\n\n    def __init__(self, dataloader: Union[List[DataLoader], DataLoader]):\n        self.dataloader = dataloader\n\n        # cannot pickle __code__ so cannot verify if PatchDataloader\n        # exists which shows dataloader methods have been overwritten.\n        # so, we hack it by using the string representation\n        self.patch_loader_code = str(self.__call__.__code__)\n\n    def __call__(self) -> Union[List[DataLoader], DataLoader]:\n        return self.dataloader\n", "path": "pytorch_lightning/trainer/connectors/data_connector.py"}]}