To create a confusion matrix you can use
def get_preds(model, data_loader):
    with torch.no_grad():
        preds = torch.tensor([]) # It will hold tuple (predicted_label, real_label)
        for batch in loader:
            images, labels = batch # labels are already real_labels (0,1,2, ...)

            pred = model(images) # pred.shape = [batch_size, num_classes]
            pred = torch.argmax(-1) # this will give you index of the predicted class
            preds = torch.cat((pred, labels), dim=0)
    return preds # shape is [total_images, 2]

# Now you can use a for loop to get the confusion matrix
num_classes = 4
c_mat = torch.zeros(num_classes, num_classes, dtype=torch.int64)

for p in preds:
    i,j = p[0], p[1] # Add code here to convert indices to int
    c_mat[i,j] += 1

I still don’t understand what you mean by and a single image.