Merge pull request yunjey#49 from Kongsea/master

yunjey · web-flow · commit 4b67434961a6 · 2017-07-11T13:35:01.000+09:00
Update main-gpu.py
diff --git a/tutorials/01-basics/feedforward_neural_network/main-gpu.py b/tutorials/01-basics/feedforward_neural_network/main-gpu.py
@@ -57,8 +57,8 @@ def forward(self, x):
 for epoch in range(num_epochs):
     for i, (images, labels) in enumerate(train_loader):  
         # Convert torch tensor to Variable
-        images = Variable(images.view(-1, 28*28)).cuda()
-        labels = Variable(labels).cuda()
+        images = Variable(images.view(-1, 28*28).cuda())
+        labels = Variable(labels.cuda())
         
         # Forward + Backward + Optimize
         optimizer.zero_grad()  # zero the gradient buffer
@@ -84,4 +84,4 @@ def forward(self, x):
 print('Accuracy of the network on the 10000 test images: %d %%' % (100 * correct / total))
 
 # Save the Model
-torch.save(net.state_dict(), 'model.pkl')
+torch.save(net.state_dict(), 'model.pkl')
