CNN Predictions

Hi everyone,

In this example:

for images, labels in test_loader:
    images = Variable(images).cuda()
    outputs = cnn(images)
    _, predicted = torch.max(, 1)

I would like to know why we define the predictions by finding the maximum between and 1.

I did a Google search and it seems like you got that code from this example:

The CNN network is trained with CrossEntropyLoss which will mean the maximum value given by the final layer will be the predicted result. The 1 given as a parameter to torch.max is simply the dimension in we want to check for the maximum:

 torch.max(input, dim, max=None, max_indices=None) -> (Tensor, LongTensor)


1 Like