Checking trained CNN on the images
torch.argmax has an extra argument dim which you can specify such that the maximum value is taken over a specific dimension. If you specify the dimension which represents the number of images it will return an array of indices where each value is for one image. For example:
import torch

# 3 images with 5 classes
t = torch.randn(3, 5)

# tensor([[-1.2917,  1.3740,  0.6967, -0.0575,  0.3702],
#        [ 0.5428,  1.0863,  0.3951,  1.8535,  1.0926],
#        [ 0.5865,  0.8522, -0.6858,  0.5297, -0.1320]])

# get the argmax over the first dimension, which specifies the number of images
torch.argmax(t, dim=1)

# tensor([1, 3, 1])