prompt: 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()
  
  ```
  
golden_diff: 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
  
verification_info: {"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"}]}