example: pytorch_lightning_mnist.py (#3290)

Set GPU device with horovod local rank.

Signed-off-by: Chongxiao Cao <chongxiaoc@uber.com>

examples/pytorch/pytorch_lightning_mnist.py

@@ -110,8 +110,12 @@ def test():
 
 if __name__ == '__main__':
     args = parser.parse_args()
-    args.cuda = not args.no_cuda and torch.cuda.is_available()
+    torch.manual_seed(args.seed)
     hvd.init()
+    args.cuda = not args.no_cuda and torch.cuda.is_available()
+    if args.cuda:
+        torch.cuda.set_device(hvd.local_rank())
+        torch.cuda.manual_seed(args.seed)
 
     kwargs = {'num_workers': 2}
     # When supported, use 'forkserver' to spawn dataloader workers instead of 'fork' to prevent
@@ -205,4 +209,3 @@ def on_train_end(self, trainer, model):
         if args.cuda:
             model = model.cuda()
         test()
-