Repository: horovod/horovod

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


Discussion:

Resolving patch:
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 @@ def on_train_end(self, trainer, model):
         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 @@ def on_train_end(self, trainer, model):
                           num_sanity_val_steps=0)
 
         trainer.fit(model)
-
+        if args.cuda:
+            model = model.cuda()
         test()
 


Test patch: