Predict test data

I have my model ready trained on data with labels called ‘labels2’ that have 0 and 1. I want to predict the new data and for it to have labels 0, 1, and 2 if it doesn’t recognize it.
This is the code I have written so far.
the dataset in imgset_loader’s shape is (4, 1, 32, 32)

labels=[0,1]
for i, images in enumerate(imgset_loader):
    images = images.to(device)
    net = net.double()
    outputs = net(images)
    _, predicted = torch.max(outputs, 1)

    print((i,labels[predicted]))

But I think that something is wrong with it because I keep getting 1 as labels for all of them and when reloading the data it would show something random.

would it also be possible to show the name of the image along with the predicted value?

Why you use only the first element of predicted?

predicted is a tensor of 1 but even if I change it to predicted[i] for example. it would still come off as 1

is your minibatch size 1?
also - how exactly the model should guess the label 2 when you teach it only on label 0 and 1?

yes, it is a batch of 1.
and about label 2 it’s actually part of my question. do I need to retrain the model to recognize label 2?

There is no such thing as “doesn’t recognize” as in:

I want to predict the new data and for it to have labels 0, 1, and 2 if it doesn’t recognize it.

There will be always either 0 prediction or 1 prediction.

In other words, the model must be trained to recognize all labels you want to predict.