Fix up test code

Youngyi · Jun 19, 2017 · 528d7cc · 528d7cc
1 parent 51f1c93
commit 528d7cc
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/tutorial-contents/401_CNN.py b/tutorial-contents/401_CNN.py
@@ -122,7 +122,7 @@ def plot_with_labels(lowDWeights, labels):
 plt.ioff()
 
 # print 10 predictions from test data
-test_output = cnn(test_x[:10])
+test_output, _ = cnn(test_x[:10])
 pred_y = torch.max(test_output, 1)[1].data.numpy().squeeze()
 print(pred_y, 'prediction number')
 print(test_y[:10].numpy(), 'real number')