Use argmax() instead of max()[1] in mnist/main.py (pytorch#494)

Thanks for great examples!
For beginners, I thought using `.argmax()` will be better than `.max()[1]`.

Author: seungwonpark

mnist/main.py

@@ -48,7 +48,7 @@ def test(args, model, device, test_loader):
             data, target = data.to(device), target.to(device)
             output = model(data)
             test_loss += F.nll_loss(output, target, reduction='sum').item() # sum up batch loss
-            pred = output.max(1, keepdim=True)[1] # get the index of the max log-probability
+            pred = output.argmax(dim=1, keepdim=True) # get the index of the max log-probability
             correct += pred.eq(target.view_as(pred)).sum().item()
 
     test_loss /= len(test_loader.dataset)