We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
pytorch_lightning_mnist.py example is not working with GPU
**Environment:**
torch==1.7.1
pytorch-lightning==1.3.8
horovod==master branch

**Bug report:**
Reproduce with `horovodrun -np 1 -H agent11240-phx4:1 python pytorch_lightning_mnist.py --epochs 1`
It fails in test step since model weights are on cpu:
```
Epoch 0: 100% 946/948 [00:10<00:00, 86.76it/s, loss=0.438, v[[1,0]<stdout>:A epoch ended.
Epoch 0: 100% 948/948 [00:11<00:00, 84.24it/s, loss=0.438, v_[1,0]<stdout>:Training ends
Epoch 0: 100% 948/948 [00:11<00:00, 84.18it/s, loss=0.438, v_num=0][1,0]<stdout>:
[1,0]<stderr>:Traceback (most recent call last):
[1,0]<stderr>:  File "pytorch_lightning_mnist.py", line 218, in <module>
[1,0]<stderr>:    test()
[1,0]<stderr>:  File "pytorch_lightning_mnist.py", line 101, in test
[1,0]<stderr>:    output = model(data)
[1,0]<stderr>:  File "/usr/lib/python3.6/site-packages/torch/nn/modules/module.py", line 727, in _call_impl
[1,0]<stderr>:    result = self.forward(*input, **kwargs)
[1,0]<stderr>:  File "pytorch_lightning_mnist.py", line 59, in forward
[1,0]<stderr>:    x = F.relu(F.max_pool2d(self.conv1(x), 2))
[1,0]<stderr>:  File "/usr/lib/python3.6/site-packages/torch/nn/modules/module.py", line 727, in _call_impl
[1,0]<stderr>:    result = self.forward(*input, **kwargs)
[1,0]<stderr>:  File "/usr/lib/python3.6/site-packages/torch/nn/modules/conv.py", line 423, in forward
[1,0]<stderr>:    return self._conv_forward(input, self.weight)
[1,0]<stderr>:  File "/usr/lib/python3.6/site-packages/torch/nn/modules/conv.py", line 420, in _conv_forward
[1,0]<stderr>:    self.padding, self.dilation, self.groups)
[1,0]<stderr>:RuntimeError: Input type (torch.cuda.FloatTensor) and weight type (torch.FloatTensor) should be the same
```

CPU version `horovodrun -np 1 -H agent11240-phx4:1 python pytorch_lightning_mnist.py --epochs 1 --no-cuda` is working.

Also, the example seems has wrong configuration for using GPUs:
The `gpus` for `trainer` should depend on `args.cuda`, instead of system default setting.
https://github.com/horovod/horovod/blob/master/examples/pytorch/pytorch_lightning_mnist.py#L208

--- END ISSUE ---

Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
--- BEGIN FILES ---
Path: `examples/pytorch/pytorch_lightning_mnist.py`
Content:
```
1 import argparse
2 import os
3 from filelock import FileLock
4 import tempfile
5 
6 import torch
7 import torch.multiprocessing as mp
8 import torch.nn as nn
9 import torch.nn.functional as F
10 import torch.optim as optim
11 from torchvision import datasets, transforms
12 # import torch.utils.data.distributed
13 
14 from pytorch_lightning import LightningModule, Trainer
15 from pytorch_lightning.callbacks import ModelCheckpoint
16 from pytorch_lightning.loggers import TensorBoardLogger
17 import horovod.torch as hvd
18 
19 # Training settings
20 parser = argparse.ArgumentParser(description='PyTorch MNIST Example')
21 parser.add_argument('--batch-size', type=int, default=64, metavar='N',
22                     help='input batch size for training (default: 64)')
23 parser.add_argument('--test-batch-size', type=int, default=1000, metavar='N',
24                     help='input batch size for testing (default: 1000)')
25 parser.add_argument('--epochs', type=int, default=10, metavar='N',
26                     help='number of epochs to train (default: 10)')
27 parser.add_argument('--lr', type=float, default=0.01, metavar='LR',
28                     help='learning rate (default: 0.01)')
29 parser.add_argument('--momentum', type=float, default=0.5, metavar='M',
30                     help='SGD momentum (default: 0.5)')
31 parser.add_argument('--no-cuda', action='store_true', default=False,
32                     help='disables CUDA training')
33 parser.add_argument('--seed', type=int, default=42, metavar='S',
34                     help='random seed (default: 42)')
35 parser.add_argument('--log-interval', type=int, default=10, metavar='N',
36                     help='how many batches to wait before logging training status')
37 parser.add_argument('--fp16-allreduce', action='store_true', default=False,
38                     help='use fp16 compression during allreduce')
39 parser.add_argument('--use-adasum', action='store_true', default=False,
40                     help='use adasum algorithm to do reduction')
41 parser.add_argument('--gradient-predivide-factor', type=float, default=1.0,
42                     help='apply gradient predivide factor in optimizer (default: 1.0)')
43 parser.add_argument('--data-dir',
44                     help='location of the training dataset in the local filesystem (will be downloaded if needed)')
45 
46 
47 # Define the PyTorch model without any Horovod-specific parameters
48 class Net(LightningModule):
49     def __init__(self):
50         super(Net, self).__init__()
51         self.conv1 = nn.Conv2d(1, 10, kernel_size=5)
52         self.conv2 = nn.Conv2d(10, 20, kernel_size=5)
53         self.conv2_drop = nn.Dropout2d()
54         self.fc1 = nn.Linear(320, 50)
55         self.fc2 = nn.Linear(50, 10)
56 
57     def forward(self, x):
58         x = x.float()
59         x = F.relu(F.max_pool2d(self.conv1(x), 2))
60         x = F.relu(F.max_pool2d(self.conv2_drop(self.conv2(x)), 2))
61         x = x.view(-1, 320)
62         x = F.relu(self.fc1(x))
63         x = F.dropout(x, training=self.training)
64         x = self.fc2(x)
65         return F.log_softmax(x, -1)
66 
67     def configure_optimizers(self):
68         return optim.SGD(self.parameters(), lr=0.01, momentum=0.5)
69 
70     def training_step(self, batch, batch_nb):
71         x, y = batch[0], batch[1]
72         y_hat = self(x)
73         loss = F.nll_loss(y_hat, y.long())
74         tensorboard_logs = {'train_loss': loss}
75         return {'loss': loss, 'log': tensorboard_logs}
76 
77     def validation_step(self, batch, batch_nb):
78         x, y = batch[0], batch[1]
79         y_hat = self(x)
80         return {'val_loss': F.nll_loss(y_hat, y.long())}
81 
82     def validation_epoch_end(self, outputs):
83         avg_loss = torch.stack([x['val_loss'] for x in outputs]).mean()
84         tensorboard_logs = {'val_loss': avg_loss}
85         return {'avg_val_loss': avg_loss, 'log': tensorboard_logs}
86 
87 
88 def metric_average(val, name):
89     tensor = torch.tensor(val)
90     avg_tensor = hvd.allreduce(tensor, name=name)
91     return avg_tensor.item()
92 
93 
94 def test():
95     model.eval()
96     test_loss = 0.
97     test_accuracy = 0.
98     for data, target in test_loader:
99         if args.cuda:
100             data, target = data.cuda(), target.cuda()
101         output = model(data)
102         # sum up batch loss
103         test_loss += F.nll_loss(output, target, size_average=False).item()
104         # get the index of the max log-probability
105         pred = output.data.max(1, keepdim=True)[1]
106         test_accuracy += pred.eq(target.data.view_as(pred)).cpu().float().sum()
107 
108     # Horovod: use test_sampler to determine the number of examples in
109     # this worker's partition.
110     test_loss /= len(test_sampler)
111     test_accuracy /= len(test_sampler)
112 
113     # Horovod: average metric values across workers.
114     test_loss = metric_average(test_loss, 'avg_loss')
115     test_accuracy = metric_average(test_accuracy, 'avg_accuracy')
116 
117     # Horovod: print output only on first rank.
118     if hvd.rank() == 0:
119         print('\nTest set: Average loss: {:.4f}, Accuracy: {:.2f}%\n'.format(
120             test_loss, 100. * test_accuracy))
121 
122 
123 if __name__ == '__main__':
124     args = parser.parse_args()
125     args.cuda = not args.no_cuda and torch.cuda.is_available()
126     hvd.init()
127 
128     kwargs = {'num_workers': 2}
129     # When supported, use 'forkserver' to spawn dataloader workers instead of 'fork' to prevent
130     # issues with Infiniband implementations that are not fork-safe
131     if (kwargs.get('num_workers', 0) > 0 and hasattr(mp, '_supports_context') and
132             mp._supports_context and 'forkserver' in mp.get_all_start_methods()):
133         kwargs['multiprocessing_context'] = 'forkserver'
134 
135     # get data
136     data_dir = args.data_dir or './data'
137     with FileLock(os.path.expanduser("~/.horovod_lock")):
138         train_dataset = \
139             datasets.MNIST(data_dir, train=True, download=True,
140                            transform=transforms.Compose([
141                                transforms.ToTensor(),
142                                transforms.Normalize((0.1307,), (0.3081,))
143                            ]))
144 
145     # set training data loader
146     train_sampler = torch.utils.data.distributed.DistributedSampler(
147         train_dataset, num_replicas=hvd.size(), rank=hvd.rank())
148     train_loader = torch.utils.data.DataLoader(
149         train_dataset, batch_size=args.batch_size, sampler=train_sampler, **kwargs)
150 
151     test_dataset = \
152         datasets.MNIST(data_dir, train=False, transform=transforms.Compose([
153             transforms.ToTensor(),
154             transforms.Normalize((0.1307,), (0.3081,))
155         ]))
156 
157     # set validation data loader
158     test_sampler = torch.utils.data.distributed.DistributedSampler(
159         test_dataset, num_replicas=hvd.size(), rank=hvd.rank())
160     test_loader = torch.utils.data.DataLoader(test_dataset, batch_size=args.test_batch_size,
161                                               sampler=test_sampler, **kwargs)
162 
163     epochs = args.epochs
164     with tempfile.TemporaryDirectory() as run_output_dir:
165         ckpt_path = os.path.join(run_output_dir, "checkpoint")
166         os.makedirs(ckpt_path, exist_ok=True)
167 
168         logs_path = os.path.join(run_output_dir, "logger")
169         os.makedirs(logs_path, exist_ok=True)
170         logger = TensorBoardLogger(logs_path)
171 
172         train_percent = 1.0
173         val_percent = 1.0
174 
175         model = Net()
176         setattr(model, 'train_dataloader', lambda: train_loader)
177         setattr(model, 'val_dataloader', lambda: test_loader)
178 
179         from pytorch_lightning.callbacks import Callback
180 
181         class MyDummyCallback(Callback):
182             def __init__(self):
183                 self.epcoh_end_counter = 0
184                 self.train_epcoh_end_counter = 0
185 
186             def on_init_start(self, trainer):
187                 print('Starting to init trainer!')
188 
189             def on_init_end(self, trainer):
190                 print('Trainer is initialized.')
191 
192             def on_epoch_end(self, trainer, model):
193                 print('A epoch ended.')
194                 self.epcoh_end_counter += 1
195 
196             def on_train_epoch_end(self, trainer, model, unused=None):
197                 print('A train epoch ended.')
198                 self.train_epcoh_end_counter += 1
199 
200             def on_train_end(self, trainer, model):
201                 print('Training ends')
202                 assert self.epcoh_end_counter == 2 * epochs
203                 assert self.train_epcoh_end_counter == epochs
204 
205         callbacks = [MyDummyCallback(), ModelCheckpoint(dirpath=ckpt_path)]
206 
207         trainer = Trainer(accelerator='horovod',
208                           gpus=(1 if torch.cuda.is_available() else 0),
209                           callbacks=callbacks,
210                           max_epochs=epochs,
211                           limit_train_batches=train_percent,
212                           limit_val_batches=val_percent,
213                           logger=logger,
214                           num_sanity_val_steps=0)
215 
216         trainer.fit(model)
217 
218         test()
219 
220
```
--- 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/examples/pytorch/pytorch_lightning_mnist.py b/examples/pytorch/pytorch_lightning_mnist.py
--- a/examples/pytorch/pytorch_lightning_mnist.py
+++ b/examples/pytorch/pytorch_lightning_mnist.py
@@ -24,22 +24,10 @@
                     help='input batch size for testing (default: 1000)')
 parser.add_argument('--epochs', type=int, default=10, metavar='N',
                     help='number of epochs to train (default: 10)')
-parser.add_argument('--lr', type=float, default=0.01, metavar='LR',
-                    help='learning rate (default: 0.01)')
-parser.add_argument('--momentum', type=float, default=0.5, metavar='M',
-                    help='SGD momentum (default: 0.5)')
 parser.add_argument('--no-cuda', action='store_true', default=False,
                     help='disables CUDA training')
 parser.add_argument('--seed', type=int, default=42, metavar='S',
                     help='random seed (default: 42)')
-parser.add_argument('--log-interval', type=int, default=10, metavar='N',
-                    help='how many batches to wait before logging training status')
-parser.add_argument('--fp16-allreduce', action='store_true', default=False,
-                    help='use fp16 compression during allreduce')
-parser.add_argument('--use-adasum', action='store_true', default=False,
-                    help='use adasum algorithm to do reduction')
-parser.add_argument('--gradient-predivide-factor', type=float, default=1.0,
-                    help='apply gradient predivide factor in optimizer (default: 1.0)')
 parser.add_argument('--data-dir',
                     help='location of the training dataset in the local filesystem (will be downloaded if needed)')
 
@@ -205,7 +193,7 @@
         callbacks = [MyDummyCallback(), ModelCheckpoint(dirpath=ckpt_path)]
 
         trainer = Trainer(accelerator='horovod',
-                          gpus=(1 if torch.cuda.is_available() else 0),
+                          gpus=(1 if args.cuda else 0),
                           callbacks=callbacks,
                           max_epochs=epochs,
                           limit_train_batches=train_percent,
@@ -214,6 +202,7 @@
                           num_sanity_val_steps=0)
 
         trainer.fit(model)
-
+        if args.cuda:
+            model = model.cuda()
         test()

{"golden_diff": "diff --git a/examples/pytorch/pytorch_lightning_mnist.py b/examples/pytorch/pytorch_lightning_mnist.py\n--- a/examples/pytorch/pytorch_lightning_mnist.py\n+++ b/examples/pytorch/pytorch_lightning_mnist.py\n@@ -24,22 +24,10 @@\n                     help='input batch size for testing (default: 1000)')\n parser.add_argument('--epochs', type=int, default=10, metavar='N',\n                     help='number of epochs to train (default: 10)')\n-parser.add_argument('--lr', type=float, default=0.01, metavar='LR',\n-                    help='learning rate (default: 0.01)')\n-parser.add_argument('--momentum', type=float, default=0.5, metavar='M',\n-                    help='SGD momentum (default: 0.5)')\n parser.add_argument('--no-cuda', action='store_true', default=False,\n                     help='disables CUDA training')\n parser.add_argument('--seed', type=int, default=42, metavar='S',\n                     help='random seed (default: 42)')\n-parser.add_argument('--log-interval', type=int, default=10, metavar='N',\n-                    help='how many batches to wait before logging training status')\n-parser.add_argument('--fp16-allreduce', action='store_true', default=False,\n-                    help='use fp16 compression during allreduce')\n-parser.add_argument('--use-adasum', action='store_true', default=False,\n-                    help='use adasum algorithm to do reduction')\n-parser.add_argument('--gradient-predivide-factor', type=float, default=1.0,\n-                    help='apply gradient predivide factor in optimizer (default: 1.0)')\n parser.add_argument('--data-dir',\n                     help='location of the training dataset in the local filesystem (will be downloaded if needed)')\n \n@@ -205,7 +193,7 @@\n         callbacks = [MyDummyCallback(), ModelCheckpoint(dirpath=ckpt_path)]\n \n         trainer = Trainer(accelerator='horovod',\n-                          gpus=(1 if torch.cuda.is_available() else 0),\n+                          gpus=(1 if args.cuda else 0),\n                           callbacks=callbacks,\n                           max_epochs=epochs,\n                           limit_train_batches=train_percent,\n@@ -214,6 +202,7 @@\n                           num_sanity_val_steps=0)\n \n         trainer.fit(model)\n-\n+        if args.cuda:\n+            model = model.cuda()\n         test()\n", "issue": "pytorch_lightning_mnist.py example is not working with GPU\n**Environment:**\r\ntorch==1.7.1\r\npytorch-lightning==1.3.8\r\nhorovod==master branch\r\n\r\n**Bug report:**\r\nReproduce with `horovodrun -np 1 -H agent11240-phx4:1 python pytorch_lightning_mnist.py --epochs 1`\r\nIt fails in test step since model weights are on cpu:\r\n```\r\nEpoch 0: 100% 946/948 [00:10<00:00, 86.76it/s, loss=0.438, v[[1,0]<stdout>:A epoch ended.\r\nEpoch 0: 100% 948/948 [00:11<00:00, 84.24it/s, loss=0.438, v_[1,0]<stdout>:Training ends\r\nEpoch 0: 100% 948/948 [00:11<00:00, 84.18it/s, loss=0.438, v_num=0][1,0]<stdout>:\r\n[1,0]<stderr>:Traceback (most recent call last):\r\n[1,0]<stderr>:  File \"pytorch_lightning_mnist.py\", line 218, in <module>\r\n[1,0]<stderr>:    test()\r\n[1,0]<stderr>:  File \"pytorch_lightning_mnist.py\", line 101, in test\r\n[1,0]<stderr>:    output = model(data)\r\n[1,0]<stderr>:  File \"/usr/lib/python3.6/site-packages/torch/nn/modules/module.py\", line 727, in _call_impl\r\n[1,0]<stderr>:    result = self.forward(*input, **kwargs)\r\n[1,0]<stderr>:  File \"pytorch_lightning_mnist.py\", line 59, in forward\r\n[1,0]<stderr>:    x = F.relu(F.max_pool2d(self.conv1(x), 2))\r\n[1,0]<stderr>:  File \"/usr/lib/python3.6/site-packages/torch/nn/modules/module.py\", line 727, in _call_impl\r\n[1,0]<stderr>:    result = self.forward(*input, **kwargs)\r\n[1,0]<stderr>:  File \"/usr/lib/python3.6/site-packages/torch/nn/modules/conv.py\", line 423, in forward\r\n[1,0]<stderr>:    return self._conv_forward(input, self.weight)\r\n[1,0]<stderr>:  File \"/usr/lib/python3.6/site-packages/torch/nn/modules/conv.py\", line 420, in _conv_forward\r\n[1,0]<stderr>:    self.padding, self.dilation, self.groups)\r\n[1,0]<stderr>:RuntimeError: Input type (torch.cuda.FloatTensor) and weight type (torch.FloatTensor) should be the same\r\n```\r\n\r\nCPU version `horovodrun -np 1 -H agent11240-phx4:1 python pytorch_lightning_mnist.py --epochs 1 --no-cuda` is working.\r\n\r\nAlso, the example seems has wrong configuration for using GPUs:\r\nThe `gpus` for `trainer` should depend on `args.cuda`, instead of system default setting.\r\nhttps://github.com/horovod/horovod/blob/master/examples/pytorch/pytorch_lightning_mnist.py#L208\n", "before_files": [{"content": "import argparse\nimport os\nfrom filelock import FileLock\nimport tempfile\n\nimport torch\nimport torch.multiprocessing as mp\nimport torch.nn as nn\nimport torch.nn.functional as F\nimport torch.optim as optim\nfrom torchvision import datasets, transforms\n# import torch.utils.data.distributed\n\nfrom pytorch_lightning import LightningModule, Trainer\nfrom pytorch_lightning.callbacks import ModelCheckpoint\nfrom pytorch_lightning.loggers import TensorBoardLogger\nimport horovod.torch as hvd\n\n# Training settings\nparser = argparse.ArgumentParser(description='PyTorch MNIST Example')\nparser.add_argument('--batch-size', type=int, default=64, metavar='N',\n                    help='input batch size for training (default: 64)')\nparser.add_argument('--test-batch-size', type=int, default=1000, metavar='N',\n                    help='input batch size for testing (default: 1000)')\nparser.add_argument('--epochs', type=int, default=10, metavar='N',\n                    help='number of epochs to train (default: 10)')\nparser.add_argument('--lr', type=float, default=0.01, metavar='LR',\n                    help='learning rate (default: 0.01)')\nparser.add_argument('--momentum', type=float, default=0.5, metavar='M',\n                    help='SGD momentum (default: 0.5)')\nparser.add_argument('--no-cuda', action='store_true', default=False,\n                    help='disables CUDA training')\nparser.add_argument('--seed', type=int, default=42, metavar='S',\n                    help='random seed (default: 42)')\nparser.add_argument('--log-interval', type=int, default=10, metavar='N',\n                    help='how many batches to wait before logging training status')\nparser.add_argument('--fp16-allreduce', action='store_true', default=False,\n                    help='use fp16 compression during allreduce')\nparser.add_argument('--use-adasum', action='store_true', default=False,\n                    help='use adasum algorithm to do reduction')\nparser.add_argument('--gradient-predivide-factor', type=float, default=1.0,\n                    help='apply gradient predivide factor in optimizer (default: 1.0)')\nparser.add_argument('--data-dir',\n                    help='location of the training dataset in the local filesystem (will be downloaded if needed)')\n\n\n# Define the PyTorch model without any Horovod-specific parameters\nclass Net(LightningModule):\n    def __init__(self):\n        super(Net, self).__init__()\n        self.conv1 = nn.Conv2d(1, 10, kernel_size=5)\n        self.conv2 = nn.Conv2d(10, 20, kernel_size=5)\n        self.conv2_drop = nn.Dropout2d()\n        self.fc1 = nn.Linear(320, 50)\n        self.fc2 = nn.Linear(50, 10)\n\n    def forward(self, x):\n        x = x.float()\n        x = F.relu(F.max_pool2d(self.conv1(x), 2))\n        x = F.relu(F.max_pool2d(self.conv2_drop(self.conv2(x)), 2))\n        x = x.view(-1, 320)\n        x = F.relu(self.fc1(x))\n        x = F.dropout(x, training=self.training)\n        x = self.fc2(x)\n        return F.log_softmax(x, -1)\n\n    def configure_optimizers(self):\n        return optim.SGD(self.parameters(), lr=0.01, momentum=0.5)\n\n    def training_step(self, batch, batch_nb):\n        x, y = batch[0], batch[1]\n        y_hat = self(x)\n        loss = F.nll_loss(y_hat, y.long())\n        tensorboard_logs = {'train_loss': loss}\n        return {'loss': loss, 'log': tensorboard_logs}\n\n    def validation_step(self, batch, batch_nb):\n        x, y = batch[0], batch[1]\n        y_hat = self(x)\n        return {'val_loss': F.nll_loss(y_hat, y.long())}\n\n    def validation_epoch_end(self, outputs):\n        avg_loss = torch.stack([x['val_loss'] for x in outputs]).mean()\n        tensorboard_logs = {'val_loss': avg_loss}\n        return {'avg_val_loss': avg_loss, 'log': tensorboard_logs}\n\n\ndef metric_average(val, name):\n    tensor = torch.tensor(val)\n    avg_tensor = hvd.allreduce(tensor, name=name)\n    return avg_tensor.item()\n\n\ndef test():\n    model.eval()\n    test_loss = 0.\n    test_accuracy = 0.\n    for data, target in test_loader:\n        if args.cuda:\n            data, target = data.cuda(), target.cuda()\n        output = model(data)\n        # sum up batch loss\n        test_loss += F.nll_loss(output, target, size_average=False).item()\n        # get the index of the max log-probability\n        pred = output.data.max(1, keepdim=True)[1]\n        test_accuracy += pred.eq(target.data.view_as(pred)).cpu().float().sum()\n\n    # Horovod: use test_sampler to determine the number of examples in\n    # this worker's partition.\n    test_loss /= len(test_sampler)\n    test_accuracy /= len(test_sampler)\n\n    # Horovod: average metric values across workers.\n    test_loss = metric_average(test_loss, 'avg_loss')\n    test_accuracy = metric_average(test_accuracy, 'avg_accuracy')\n\n    # Horovod: print output only on first rank.\n    if hvd.rank() == 0:\n        print('\\nTest set: Average loss: {:.4f}, Accuracy: {:.2f}%\\n'.format(\n            test_loss, 100. * test_accuracy))\n\n\nif __name__ == '__main__':\n    args = parser.parse_args()\n    args.cuda = not args.no_cuda and torch.cuda.is_available()\n    hvd.init()\n\n    kwargs = {'num_workers': 2}\n    # When supported, use 'forkserver' to spawn dataloader workers instead of 'fork' to prevent\n    # issues with Infiniband implementations that are not fork-safe\n    if (kwargs.get('num_workers', 0) > 0 and hasattr(mp, '_supports_context') and\n            mp._supports_context and 'forkserver' in mp.get_all_start_methods()):\n        kwargs['multiprocessing_context'] = 'forkserver'\n\n    # get data\n    data_dir = args.data_dir or './data'\n    with FileLock(os.path.expanduser(\"~/.horovod_lock\")):\n        train_dataset = \\\n            datasets.MNIST(data_dir, train=True, download=True,\n                           transform=transforms.Compose([\n                               transforms.ToTensor(),\n                               transforms.Normalize((0.1307,), (0.3081,))\n                           ]))\n\n    # set training data loader\n    train_sampler = torch.utils.data.distributed.DistributedSampler(\n        train_dataset, num_replicas=hvd.size(), rank=hvd.rank())\n    train_loader = torch.utils.data.DataLoader(\n        train_dataset, batch_size=args.batch_size, sampler=train_sampler, **kwargs)\n\n    test_dataset = \\\n        datasets.MNIST(data_dir, train=False, transform=transforms.Compose([\n            transforms.ToTensor(),\n            transforms.Normalize((0.1307,), (0.3081,))\n        ]))\n\n    # set validation data loader\n    test_sampler = torch.utils.data.distributed.DistributedSampler(\n        test_dataset, num_replicas=hvd.size(), rank=hvd.rank())\n    test_loader = torch.utils.data.DataLoader(test_dataset, batch_size=args.test_batch_size,\n                                              sampler=test_sampler, **kwargs)\n\n    epochs = args.epochs\n    with tempfile.TemporaryDirectory() as run_output_dir:\n        ckpt_path = os.path.join(run_output_dir, \"checkpoint\")\n        os.makedirs(ckpt_path, exist_ok=True)\n\n        logs_path = os.path.join(run_output_dir, \"logger\")\n        os.makedirs(logs_path, exist_ok=True)\n        logger = TensorBoardLogger(logs_path)\n\n        train_percent = 1.0\n        val_percent = 1.0\n\n        model = Net()\n        setattr(model, 'train_dataloader', lambda: train_loader)\n        setattr(model, 'val_dataloader', lambda: test_loader)\n\n        from pytorch_lightning.callbacks import Callback\n\n        class MyDummyCallback(Callback):\n            def __init__(self):\n                self.epcoh_end_counter = 0\n                self.train_epcoh_end_counter = 0\n\n            def on_init_start(self, trainer):\n                print('Starting to init trainer!')\n\n            def on_init_end(self, trainer):\n                print('Trainer is initialized.')\n\n            def on_epoch_end(self, trainer, model):\n                print('A epoch ended.')\n                self.epcoh_end_counter += 1\n\n            def on_train_epoch_end(self, trainer, model, unused=None):\n                print('A train epoch ended.')\n                self.train_epcoh_end_counter += 1\n\n            def on_train_end(self, trainer, model):\n                print('Training ends')\n                assert self.epcoh_end_counter == 2 * epochs\n                assert self.train_epcoh_end_counter == epochs\n\n        callbacks = [MyDummyCallback(), ModelCheckpoint(dirpath=ckpt_path)]\n\n        trainer = Trainer(accelerator='horovod',\n                          gpus=(1 if torch.cuda.is_available() else 0),\n                          callbacks=callbacks,\n                          max_epochs=epochs,\n                          limit_train_batches=train_percent,\n                          limit_val_batches=val_percent,\n                          logger=logger,\n                          num_sanity_val_steps=0)\n\n        trainer.fit(model)\n\n        test()\n\n", "path": "examples/pytorch/pytorch_lightning_mnist.py"}], "after_files": [{"content": "import argparse\nimport os\nfrom filelock import FileLock\nimport tempfile\n\nimport torch\nimport torch.multiprocessing as mp\nimport torch.nn as nn\nimport torch.nn.functional as F\nimport torch.optim as optim\nfrom torchvision import datasets, transforms\n# import torch.utils.data.distributed\n\nfrom pytorch_lightning import LightningModule, Trainer\nfrom pytorch_lightning.callbacks import ModelCheckpoint\nfrom pytorch_lightning.loggers import TensorBoardLogger\nimport horovod.torch as hvd\n\n# Training settings\nparser = argparse.ArgumentParser(description='PyTorch MNIST Example')\nparser.add_argument('--batch-size', type=int, default=64, metavar='N',\n                    help='input batch size for training (default: 64)')\nparser.add_argument('--test-batch-size', type=int, default=1000, metavar='N',\n                    help='input batch size for testing (default: 1000)')\nparser.add_argument('--epochs', type=int, default=10, metavar='N',\n                    help='number of epochs to train (default: 10)')\nparser.add_argument('--no-cuda', action='store_true', default=False,\n                    help='disables CUDA training')\nparser.add_argument('--seed', type=int, default=42, metavar='S',\n                    help='random seed (default: 42)')\nparser.add_argument('--data-dir',\n                    help='location of the training dataset in the local filesystem (will be downloaded if needed)')\n\n\n# Define the PyTorch model without any Horovod-specific parameters\nclass Net(LightningModule):\n    def __init__(self):\n        super(Net, self).__init__()\n        self.conv1 = nn.Conv2d(1, 10, kernel_size=5)\n        self.conv2 = nn.Conv2d(10, 20, kernel_size=5)\n        self.conv2_drop = nn.Dropout2d()\n        self.fc1 = nn.Linear(320, 50)\n        self.fc2 = nn.Linear(50, 10)\n\n    def forward(self, x):\n        x = x.float()\n        x = F.relu(F.max_pool2d(self.conv1(x), 2))\n        x = F.relu(F.max_pool2d(self.conv2_drop(self.conv2(x)), 2))\n        x = x.view(-1, 320)\n        x = F.relu(self.fc1(x))\n        x = F.dropout(x, training=self.training)\n        x = self.fc2(x)\n        return F.log_softmax(x, -1)\n\n    def configure_optimizers(self):\n        return optim.SGD(self.parameters(), lr=0.01, momentum=0.5)\n\n    def training_step(self, batch, batch_nb):\n        x, y = batch[0], batch[1]\n        y_hat = self(x)\n        loss = F.nll_loss(y_hat, y.long())\n        tensorboard_logs = {'train_loss': loss}\n        return {'loss': loss, 'log': tensorboard_logs}\n\n    def validation_step(self, batch, batch_nb):\n        x, y = batch[0], batch[1]\n        y_hat = self(x)\n        return {'val_loss': F.nll_loss(y_hat, y.long())}\n\n    def validation_epoch_end(self, outputs):\n        avg_loss = torch.stack([x['val_loss'] for x in outputs]).mean()\n        tensorboard_logs = {'val_loss': avg_loss}\n        return {'avg_val_loss': avg_loss, 'log': tensorboard_logs}\n\n\ndef metric_average(val, name):\n    tensor = torch.tensor(val)\n    avg_tensor = hvd.allreduce(tensor, name=name)\n    return avg_tensor.item()\n\n\ndef test():\n    model.eval()\n    test_loss = 0.\n    test_accuracy = 0.\n    for data, target in test_loader:\n        if args.cuda:\n            data, target = data.cuda(), target.cuda()\n        output = model(data)\n        # sum up batch loss\n        test_loss += F.nll_loss(output, target, size_average=False).item()\n        # get the index of the max log-probability\n        pred = output.data.max(1, keepdim=True)[1]\n        test_accuracy += pred.eq(target.data.view_as(pred)).cpu().float().sum()\n\n    # Horovod: use test_sampler to determine the number of examples in\n    # this worker's partition.\n    test_loss /= len(test_sampler)\n    test_accuracy /= len(test_sampler)\n\n    # Horovod: average metric values across workers.\n    test_loss = metric_average(test_loss, 'avg_loss')\n    test_accuracy = metric_average(test_accuracy, 'avg_accuracy')\n\n    # Horovod: print output only on first rank.\n    if hvd.rank() == 0:\n        print('\\nTest set: Average loss: {:.4f}, Accuracy: {:.2f}%\\n'.format(\n            test_loss, 100. * test_accuracy))\n\n\nif __name__ == '__main__':\n    args = parser.parse_args()\n    args.cuda = not args.no_cuda and torch.cuda.is_available()\n    hvd.init()\n\n    kwargs = {'num_workers': 2}\n    # When supported, use 'forkserver' to spawn dataloader workers instead of 'fork' to prevent\n    # issues with Infiniband implementations that are not fork-safe\n    if (kwargs.get('num_workers', 0) > 0 and hasattr(mp, '_supports_context') and\n            mp._supports_context and 'forkserver' in mp.get_all_start_methods()):\n        kwargs['multiprocessing_context'] = 'forkserver'\n\n    # get data\n    data_dir = args.data_dir or './data'\n    with FileLock(os.path.expanduser(\"~/.horovod_lock\")):\n        train_dataset = \\\n            datasets.MNIST(data_dir, train=True, download=True,\n                           transform=transforms.Compose([\n                               transforms.ToTensor(),\n                               transforms.Normalize((0.1307,), (0.3081,))\n                           ]))\n\n    # set training data loader\n    train_sampler = torch.utils.data.distributed.DistributedSampler(\n        train_dataset, num_replicas=hvd.size(), rank=hvd.rank())\n    train_loader = torch.utils.data.DataLoader(\n        train_dataset, batch_size=args.batch_size, sampler=train_sampler, **kwargs)\n\n    test_dataset = \\\n        datasets.MNIST(data_dir, train=False, transform=transforms.Compose([\n            transforms.ToTensor(),\n            transforms.Normalize((0.1307,), (0.3081,))\n        ]))\n\n    # set validation data loader\n    test_sampler = torch.utils.data.distributed.DistributedSampler(\n        test_dataset, num_replicas=hvd.size(), rank=hvd.rank())\n    test_loader = torch.utils.data.DataLoader(test_dataset, batch_size=args.test_batch_size,\n                                              sampler=test_sampler, **kwargs)\n\n    epochs = args.epochs\n    with tempfile.TemporaryDirectory() as run_output_dir:\n        ckpt_path = os.path.join(run_output_dir, \"checkpoint\")\n        os.makedirs(ckpt_path, exist_ok=True)\n\n        logs_path = os.path.join(run_output_dir, \"logger\")\n        os.makedirs(logs_path, exist_ok=True)\n        logger = TensorBoardLogger(logs_path)\n\n        train_percent = 1.0\n        val_percent = 1.0\n\n        model = Net()\n        setattr(model, 'train_dataloader', lambda: train_loader)\n        setattr(model, 'val_dataloader', lambda: test_loader)\n\n        from pytorch_lightning.callbacks import Callback\n\n        class MyDummyCallback(Callback):\n            def __init__(self):\n                self.epcoh_end_counter = 0\n                self.train_epcoh_end_counter = 0\n\n            def on_init_start(self, trainer):\n                print('Starting to init trainer!')\n\n            def on_init_end(self, trainer):\n                print('Trainer is initialized.')\n\n            def on_epoch_end(self, trainer, model):\n                print('A epoch ended.')\n                self.epcoh_end_counter += 1\n\n            def on_train_epoch_end(self, trainer, model, unused=None):\n                print('A train epoch ended.')\n                self.train_epcoh_end_counter += 1\n\n            def on_train_end(self, trainer, model):\n                print('Training ends')\n                assert self.epcoh_end_counter == 2 * epochs\n                assert self.train_epcoh_end_counter == epochs\n\n        callbacks = [MyDummyCallback(), ModelCheckpoint(dirpath=ckpt_path)]\n\n        trainer = Trainer(accelerator='horovod',\n                          gpus=(1 if args.cuda else 0),\n                          callbacks=callbacks,\n                          max_epochs=epochs,\n                          limit_train_batches=train_percent,\n                          limit_val_batches=val_percent,\n                          logger=logger,\n                          num_sanity_val_steps=0)\n\n        trainer.fit(model)\n        if args.cuda:\n            model = model.cuda()\n        test()\n\n", "path": "examples/pytorch/pytorch_lightning_mnist.py"}]}