This commit is contained in:
Morvan Zhou
2017-05-24 13:30:46 +10:00
parent ca6e7431e9
commit b2922fd3f3

View File

@ -69,7 +69,7 @@ for epoch in range(EPOCH):
test_output = cnn(test_x)
# !!!!!!!! Change in here !!!!!!!!! #
pred_y = torch.max(test_output, 1)[1].cuda().data.squeeze() # Move to CPU
pred_y = torch.max(test_output, 1)[1].cuda().data.squeeze() # move the computation in GPU
accuracy = sum(pred_y == test_y) / test_y.size(0)
print('Epoch: ', epoch, '| train loss: %.4f' % loss.data[0], '| test accuracy: %.2f' % accuracy)
@ -78,7 +78,7 @@ for epoch in range(EPOCH):
test_output = cnn(test_x[:10])
# !!!!!!!! Change in here !!!!!!!!! #
pred_y = torch.max(test_output, 1)[1].cuda().data.squeeze() # Move to CPU
pred_y = torch.max(test_output, 1)[1].cuda().data.squeeze() # move the computation in GPU
print(pred_y, 'prediction number')
print(test_y[:10], 'real number')