We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
add caching for used dataset
## 🚀 Feature

atm for each CI run the datasets are downloaded (seems so as I observed som downloading failers)

### Motivation

Speed-up runs and avoids random connection failers


--- END ISSUE ---

Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
--- BEGIN FILES ---
Path: `pl_bolts/models/mnist_module.py`
Content:
```
1 import os
2 from argparse import ArgumentParser
3 from warnings import warn
4 
5 import torch
6 from pytorch_lightning import LightningModule, Trainer
7 from torch.nn import functional as F
8 from torch.utils.data import DataLoader, random_split
9 
10 try:
11     from torchvision import transforms
12     from torchvision.datasets import MNIST
13 except ModuleNotFoundError:
14     warn('You want to use `torchvision` which is not installed yet,'  # pragma: no-cover
15          ' install it with `pip install torchvision`.')
16 
17 
18 class LitMNIST(LightningModule):
19     def __init__(self, hidden_dim=128, learning_rate=1e-3, batch_size=32, num_workers=4, data_dir='', **kwargs):
20         super().__init__()
21         self.save_hyperparameters()
22 
23         self.l1 = torch.nn.Linear(28 * 28, self.hparams.hidden_dim)
24         self.l2 = torch.nn.Linear(self.hparams.hidden_dim, 10)
25 
26         self.mnist_train = None
27         self.mnist_val = None
28 
29     def forward(self, x):
30         x = x.view(x.size(0), -1)
31         x = torch.relu(self.l1(x))
32         x = torch.relu(self.l2(x))
33         return x
34 
35     def training_step(self, batch, batch_idx):
36         x, y = batch
37         y_hat = self(x)
38         loss = F.cross_entropy(y_hat, y)
39         self.log('train_loss', loss)
40         return loss
41 
42     def validation_step(self, batch, batch_idx):
43         x, y = batch
44         y_hat = self(x)
45         loss = F.cross_entropy(y_hat, y)
46         self.log('val_loss', loss)
47 
48     def test_step(self, batch, batch_idx):
49         x, y = batch
50         y_hat = self(x)
51         loss = F.cross_entropy(y_hat, y)
52         self.log('test_loss', loss)
53 
54     def configure_optimizers(self):
55         return torch.optim.Adam(self.parameters(), lr=self.hparams.learning_rate)
56 
57     def prepare_data(self):
58         MNIST(self.hparams.data_dir, train=True, download=True, transform=transforms.ToTensor())
59 
60     def train_dataloader(self):
61         dataset = MNIST(self.hparams.data_dir, train=True, download=False, transform=transforms.ToTensor())
62         mnist_train, _ = random_split(dataset, [55000, 5000])
63         loader = DataLoader(mnist_train, batch_size=self.hparams.batch_size, num_workers=self.hparams.num_workers)
64         return loader
65 
66     def val_dataloader(self):
67         dataset = MNIST(self.hparams.data_dir, train=True, download=False, transform=transforms.ToTensor())
68         _, mnist_val = random_split(dataset, [55000, 5000])
69         loader = DataLoader(mnist_val, batch_size=self.hparams.batch_size, num_workers=self.hparams.num_workers)
70         return loader
71 
72     def test_dataloader(self):
73         test_dataset = MNIST(os.getcwd(), train=False, download=True, transform=transforms.ToTensor())
74         loader = DataLoader(test_dataset, batch_size=self.hparams.batch_size, num_workers=self.hparams.num_workers)
75         return loader
76 
77     @staticmethod
78     def add_model_specific_args(parent_parser):
79         parser = ArgumentParser(parents=[parent_parser], add_help=False)
80         parser.add_argument('--batch_size', type=int, default=32)
81         parser.add_argument('--num_workers', type=int, default=4)
82         parser.add_argument('--hidden_dim', type=int, default=128)
83         parser.add_argument('--data_dir', type=str, default='')
84         parser.add_argument('--learning_rate', type=float, default=0.0001)
85         return parser
86 
87 
88 def cli_main():
89     # args
90     parser = ArgumentParser()
91     parser = Trainer.add_argparse_args(parser)
92     parser = LitMNIST.add_model_specific_args(parser)
93     args = parser.parse_args()
94 
95     # model
96     model = LitMNIST(**vars(args))
97 
98     # training
99     trainer = Trainer.from_argparse_args(args)
100     trainer.fit(model)
101 
102 
103 if __name__ == '__main__':  # pragma: no cover
104     cli_main()
105
```
--- 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/pl_bolts/models/mnist_module.py b/pl_bolts/models/mnist_module.py
--- a/pl_bolts/models/mnist_module.py
+++ b/pl_bolts/models/mnist_module.py
@@ -1,4 +1,3 @@
-import os
 from argparse import ArgumentParser
 from warnings import warn
 
@@ -70,7 +69,7 @@
         return loader
 
     def test_dataloader(self):
-        test_dataset = MNIST(os.getcwd(), train=False, download=True, transform=transforms.ToTensor())
+        test_dataset = MNIST(self.hparams.data_dir, train=False, download=True, transform=transforms.ToTensor())
         loader = DataLoader(test_dataset, batch_size=self.hparams.batch_size, num_workers=self.hparams.num_workers)
         return loader

{"golden_diff": "diff --git a/pl_bolts/models/mnist_module.py b/pl_bolts/models/mnist_module.py\n--- a/pl_bolts/models/mnist_module.py\n+++ b/pl_bolts/models/mnist_module.py\n@@ -1,4 +1,3 @@\n-import os\n from argparse import ArgumentParser\n from warnings import warn\n \n@@ -70,7 +69,7 @@\n         return loader\n \n     def test_dataloader(self):\n-        test_dataset = MNIST(os.getcwd(), train=False, download=True, transform=transforms.ToTensor())\n+        test_dataset = MNIST(self.hparams.data_dir, train=False, download=True, transform=transforms.ToTensor())\n         loader = DataLoader(test_dataset, batch_size=self.hparams.batch_size, num_workers=self.hparams.num_workers)\n         return loader\n", "issue": "add caching for used dataset\n## \ud83d\ude80 Feature\r\n\r\natm for each CI run the datasets are downloaded (seems so as I observed som downloading failers)\r\n\r\n### Motivation\r\n\r\nSpeed-up runs and avoids random connection failers\r\n\n", "before_files": [{"content": "import os\nfrom argparse import ArgumentParser\nfrom warnings import warn\n\nimport torch\nfrom pytorch_lightning import LightningModule, Trainer\nfrom torch.nn import functional as F\nfrom torch.utils.data import DataLoader, random_split\n\ntry:\n    from torchvision import transforms\n    from torchvision.datasets import MNIST\nexcept ModuleNotFoundError:\n    warn('You want to use `torchvision` which is not installed yet,'  # pragma: no-cover\n         ' install it with `pip install torchvision`.')\n\n\nclass LitMNIST(LightningModule):\n    def __init__(self, hidden_dim=128, learning_rate=1e-3, batch_size=32, num_workers=4, data_dir='', **kwargs):\n        super().__init__()\n        self.save_hyperparameters()\n\n        self.l1 = torch.nn.Linear(28 * 28, self.hparams.hidden_dim)\n        self.l2 = torch.nn.Linear(self.hparams.hidden_dim, 10)\n\n        self.mnist_train = None\n        self.mnist_val = None\n\n    def forward(self, x):\n        x = x.view(x.size(0), -1)\n        x = torch.relu(self.l1(x))\n        x = torch.relu(self.l2(x))\n        return x\n\n    def training_step(self, batch, batch_idx):\n        x, y = batch\n        y_hat = self(x)\n        loss = F.cross_entropy(y_hat, y)\n        self.log('train_loss', loss)\n        return loss\n\n    def validation_step(self, batch, batch_idx):\n        x, y = batch\n        y_hat = self(x)\n        loss = F.cross_entropy(y_hat, y)\n        self.log('val_loss', loss)\n\n    def test_step(self, batch, batch_idx):\n        x, y = batch\n        y_hat = self(x)\n        loss = F.cross_entropy(y_hat, y)\n        self.log('test_loss', loss)\n\n    def configure_optimizers(self):\n        return torch.optim.Adam(self.parameters(), lr=self.hparams.learning_rate)\n\n    def prepare_data(self):\n        MNIST(self.hparams.data_dir, train=True, download=True, transform=transforms.ToTensor())\n\n    def train_dataloader(self):\n        dataset = MNIST(self.hparams.data_dir, train=True, download=False, transform=transforms.ToTensor())\n        mnist_train, _ = random_split(dataset, [55000, 5000])\n        loader = DataLoader(mnist_train, batch_size=self.hparams.batch_size, num_workers=self.hparams.num_workers)\n        return loader\n\n    def val_dataloader(self):\n        dataset = MNIST(self.hparams.data_dir, train=True, download=False, transform=transforms.ToTensor())\n        _, mnist_val = random_split(dataset, [55000, 5000])\n        loader = DataLoader(mnist_val, batch_size=self.hparams.batch_size, num_workers=self.hparams.num_workers)\n        return loader\n\n    def test_dataloader(self):\n        test_dataset = MNIST(os.getcwd(), train=False, download=True, transform=transforms.ToTensor())\n        loader = DataLoader(test_dataset, batch_size=self.hparams.batch_size, num_workers=self.hparams.num_workers)\n        return loader\n\n    @staticmethod\n    def add_model_specific_args(parent_parser):\n        parser = ArgumentParser(parents=[parent_parser], add_help=False)\n        parser.add_argument('--batch_size', type=int, default=32)\n        parser.add_argument('--num_workers', type=int, default=4)\n        parser.add_argument('--hidden_dim', type=int, default=128)\n        parser.add_argument('--data_dir', type=str, default='')\n        parser.add_argument('--learning_rate', type=float, default=0.0001)\n        return parser\n\n\ndef cli_main():\n    # args\n    parser = ArgumentParser()\n    parser = Trainer.add_argparse_args(parser)\n    parser = LitMNIST.add_model_specific_args(parser)\n    args = parser.parse_args()\n\n    # model\n    model = LitMNIST(**vars(args))\n\n    # training\n    trainer = Trainer.from_argparse_args(args)\n    trainer.fit(model)\n\n\nif __name__ == '__main__':  # pragma: no cover\n    cli_main()\n", "path": "pl_bolts/models/mnist_module.py"}], "after_files": [{"content": "from argparse import ArgumentParser\nfrom warnings import warn\n\nimport torch\nfrom pytorch_lightning import LightningModule, Trainer\nfrom torch.nn import functional as F\nfrom torch.utils.data import DataLoader, random_split\n\ntry:\n    from torchvision import transforms\n    from torchvision.datasets import MNIST\nexcept ModuleNotFoundError:\n    warn('You want to use `torchvision` which is not installed yet,'  # pragma: no-cover\n         ' install it with `pip install torchvision`.')\n\n\nclass LitMNIST(LightningModule):\n    def __init__(self, hidden_dim=128, learning_rate=1e-3, batch_size=32, num_workers=4, data_dir='', **kwargs):\n        super().__init__()\n        self.save_hyperparameters()\n\n        self.l1 = torch.nn.Linear(28 * 28, self.hparams.hidden_dim)\n        self.l2 = torch.nn.Linear(self.hparams.hidden_dim, 10)\n\n        self.mnist_train = None\n        self.mnist_val = None\n\n    def forward(self, x):\n        x = x.view(x.size(0), -1)\n        x = torch.relu(self.l1(x))\n        x = torch.relu(self.l2(x))\n        return x\n\n    def training_step(self, batch, batch_idx):\n        x, y = batch\n        y_hat = self(x)\n        loss = F.cross_entropy(y_hat, y)\n        self.log('train_loss', loss)\n        return loss\n\n    def validation_step(self, batch, batch_idx):\n        x, y = batch\n        y_hat = self(x)\n        loss = F.cross_entropy(y_hat, y)\n        self.log('val_loss', loss)\n\n    def test_step(self, batch, batch_idx):\n        x, y = batch\n        y_hat = self(x)\n        loss = F.cross_entropy(y_hat, y)\n        self.log('test_loss', loss)\n\n    def configure_optimizers(self):\n        return torch.optim.Adam(self.parameters(), lr=self.hparams.learning_rate)\n\n    def prepare_data(self):\n        MNIST(self.hparams.data_dir, train=True, download=True, transform=transforms.ToTensor())\n\n    def train_dataloader(self):\n        dataset = MNIST(self.hparams.data_dir, train=True, download=False, transform=transforms.ToTensor())\n        mnist_train, _ = random_split(dataset, [55000, 5000])\n        loader = DataLoader(mnist_train, batch_size=self.hparams.batch_size, num_workers=self.hparams.num_workers)\n        return loader\n\n    def val_dataloader(self):\n        dataset = MNIST(self.hparams.data_dir, train=True, download=False, transform=transforms.ToTensor())\n        _, mnist_val = random_split(dataset, [55000, 5000])\n        loader = DataLoader(mnist_val, batch_size=self.hparams.batch_size, num_workers=self.hparams.num_workers)\n        return loader\n\n    def test_dataloader(self):\n        test_dataset = MNIST(self.hparams.data_dir, train=False, download=True, transform=transforms.ToTensor())\n        loader = DataLoader(test_dataset, batch_size=self.hparams.batch_size, num_workers=self.hparams.num_workers)\n        return loader\n\n    @staticmethod\n    def add_model_specific_args(parent_parser):\n        parser = ArgumentParser(parents=[parent_parser], add_help=False)\n        parser.add_argument('--batch_size', type=int, default=32)\n        parser.add_argument('--num_workers', type=int, default=4)\n        parser.add_argument('--hidden_dim', type=int, default=128)\n        parser.add_argument('--data_dir', type=str, default='')\n        parser.add_argument('--learning_rate', type=float, default=0.0001)\n        return parser\n\n\ndef cli_main():\n    # args\n    parser = ArgumentParser()\n    parser = Trainer.add_argparse_args(parser)\n    parser = LitMNIST.add_model_specific_args(parser)\n    args = parser.parse_args()\n\n    # model\n    model = LitMNIST(**vars(args))\n\n    # training\n    trainer = Trainer.from_argparse_args(args)\n    trainer.fit(model)\n\n\nif __name__ == '__main__':  # pragma: no cover\n    cli_main()\n", "path": "pl_bolts/models/mnist_module.py"}]}