Skip to content

Commit 759aa01

Browse files
authored
Merge pull request MorvanZhou#62 from keineahnung2345/502-squeeze
502 - remove squeeze for 1-D tensor
2 parents 818b4b4 + 5d6d4f0 commit 759aa01

File tree

1 file changed

+2
-2
lines changed

1 file changed

+2
-2
lines changed

tutorial-contents/502_GPU.py

+2-2
Original file line numberDiff line numberDiff line change
@@ -68,7 +68,7 @@ def forward(self, x):
6868
test_output = cnn(test_x)
6969

7070
# !!!!!!!! Change in here !!!!!!!!! #
71-
pred_y = torch.max(test_output, 1)[1].cuda().data.squeeze() # move the computation in GPU
71+
pred_y = torch.max(test_output, 1)[1].cuda().data # move the computation in GPU
7272

7373
accuracy = torch.sum(pred_y == test_y).type(torch.FloatTensor) / test_y.size(0)
7474
print('Epoch: ', epoch, '| train loss: %.4f' % loss.data.cpu().numpy(), '| test accuracy: %.2f' % accuracy)
@@ -77,7 +77,7 @@ def forward(self, x):
7777
test_output = cnn(test_x[:10])
7878

7979
# !!!!!!!! Change in here !!!!!!!!! #
80-
pred_y = torch.max(test_output, 1)[1].cuda().data.squeeze() # move the computation in GPU
80+
pred_y = torch.max(test_output, 1)[1].cuda().data # move the computation in GPU
8181

8282
print(pred_y, 'prediction number')
8383
print(test_y[:10], 'real number')

0 commit comments

Comments
 (0)