I am using a neural network to predict the quality of the Red Wine dataset, available on UCI machine Learning, using Pytorch, and Cross Entropy Loss as loss function. This is my code:

input_size = len(input_columns)
hidden_size = 12
output_size = 6 #because there are 6 classes

#Loss function
loss_fn = F.cross_entropy

class WineQuality(nn.Module):
    def __init__(self):
        # input to hidden layer
        self.linear1 = nn.Linear(input_size, hidden_size)
        # hidden layer and output
        self.linear2 = nn.Linear(hidden_size, output_size)
    def forward(self, xb): 
        out = self.linear1(xb)
        out = F.relu(out)
        out = self.linear2(out)
        return out
    def training_step(self, batch):
        inputs, targets = batch 
        # Generate predictions
        out = self(inputs) 
        # Calcuate loss
        loss = loss_fn(out,torch.argmax(targets, dim=1))
        return loss
    def validation_step(self, batch):
        inputs, targets = batch
        # Generate predictions
        out = self(inputs)
        # Calculate loss
        loss = loss_fn(out, torch.argmax(targets, dim=1))
        return {'val_loss': loss.detach()}
    def validation_epoch_end(self, outputs):
        batch_losses = [x['val_loss'] for x in outputs]
        epoch_loss = torch.stack(batch_losses).mean()   # Combine losses
        return {'val_loss': epoch_loss.item()}
    def epoch_end(self, epoch, result, num_epochs):
        # Print result every 100th epoch
        if (epoch+1) % 100 == 0 or epoch == num_epochs-1:
            print("Epoch [{}], val_loss: {:.4f}".format(epoch+1, result['val_loss']))

model = WineQuality()

def evaluate(model, val_loader):
    outputs = [model.validation_step(batch) for batch in val_loader]
    return model.validation_epoch_end(outputs)

def fit(epochs, lr, model, train_loader, val_loader, opt_func=torch.optim.SGD):
    history = []
    optimizer = opt_func(model.parameters(), lr)
    for epoch in range(epochs):
        # Training Phase 
        for batch in train_loader:
            loss = model.training_step(batch)
        # Validation phase
        result = evaluate(model, val_loader)
        model.epoch_end(epoch, result, epochs)
    return history

loss_value = evaluate(model, valid_dl)

epochs = 1000
lr = 1e-5
history = fit(epochs, lr, model, train_loader, val_loader)

I can see that the model is good and that the loss decreases. The problem is when I have to do a prediction on an example:

def predict_single(input, target, model):
    inputs = input.unsqueeze(0)
    predictions = model(inputs)
    prediction = predictions[0].detach()
    print("Input:", input)
    print("Target:", target)
    print("Prediction:", prediction)
    return prediction

input, target = val_df[1]
prediction = predict_single(input, target, model)

This returns:

Input: tensor([0.8705, 0.3900, 2.1000, 0.0650, 4.1206, 3.3000, 0.5300, 0.2610])
Target: tensor([6.])
Prediction: tensor([ 3.6465,  0.2800, -0.4561, -1.6733, -0.6519, -0.1650])

I want to see to what are associated these logits, in the sense that I know that the highest logit is associated to the predicted class, but I want to see that class. I also applied softmax to rescale these values in a probability:

prediction = F.softmax(prediction)
output = model(input.unsqueeze(0))
_,pred = output.max(1)

And the output is the following:

tensor([0.3296, 0.1361, 0.1339, 0.1324, 0.1335, 0.1346])

I don't know what is that tensor([0]). I expect my predicted label, a value like 6.1 if the target is 6. But I am not able to obtain this.

  • $\begingroup$ Shouldn't you be using argmax instead of max? And what do you mean that you expected "6.1" as label (instead of just "6")? Also, take into account that the indexing starts at 0, so with 5 classes you would not see any 6. $\endgroup$
    – noe
    Feb 19, 2021 at 12:47
  • $\begingroup$ @noe yes, 6, not 6.1... i was thinking to my previous experiment where i used l1_loss for regression. I have not understand the indexing: i have 6 classes for wine qualities: from 3 to 8. I expect my predicted label to be a number from 3 to 8 $\endgroup$
    – CasellaJr
    Feb 19, 2021 at 13:17
  • $\begingroup$ Ahh, I see, but you are wasting computation there because classes from 0 to 2 are not used. $\endgroup$
    – noe
    Feb 19, 2021 at 13:25
  • $\begingroup$ With argmax instead of max i obtain this: ValueError: not enough values to unpack (expected 2, got 1) $\endgroup$
    – CasellaJr
    Feb 19, 2021 at 13:31
  • $\begingroup$ However, what does it change between max and argmax? I still don't have the label, but only the probability $\endgroup$
    – CasellaJr
    Feb 19, 2021 at 13:32

2 Answers 2


You should have a list of actual classes, e.g. classes = ['Superman', 'Batman', ...,'Gozilla']. The model outputs per-class logits, but without your dataset interface it's hard to say what your targets is. Since it's a multiclass problem, it should be an integer between 0 and 5. I assume the order of targets and the order of classes in classes list is the same. Then, at inference time, once you get your best = output.argmax().item(), just use classes[best] to get the class prediction.

  • $\begingroup$ Is there a way to store the classes as part of the model? $\endgroup$ Dec 3, 2022 at 2:33

If you're dealing with classification problem, then model.predict is supposed to give you logits.

outputs = net(images)
_, predicted = torch.max(outputs, 1)

for i in range(num_input):

if you have only one input then the predicted class would be as following: classes[predicted[0]]

In your case:

prediction = F.softmax(prediction)
output = model(input.unsqueeze(0))
_,pred = output.max(1)

# print class name
print(classes[pred[0]])  # <- should be added
  • $\begingroup$ I don't have defined any "classes" $\endgroup$
    – CasellaJr
    Feb 19, 2021 at 13:15
  • $\begingroup$ I have a classification problem, and the logits. So the highest logit will be associated to a class, the predicted class. How can I see this class? $\endgroup$
    – CasellaJr
    Feb 19, 2021 at 13:46
  • 1
    $\begingroup$ The output layer has 6 categories. Right? If you have already defined the 6 classes (e.g ['a', 'b', .., 'f']). So you need to give the pred output to that class classes[pred]. Have you defined 6 classes? $\endgroup$
    – yakhyo
    Feb 19, 2021 at 15:15
  • $\begingroup$ sorry i am a little bit confused: here there is all my code: jovian.ai/casella0798/problem-with-classes $\endgroup$
    – CasellaJr
    Feb 19, 2021 at 17:35
  • $\begingroup$ you can see that my target variable is "quality" that assumes 6 values: from 3 to 8 $\endgroup$
    – CasellaJr
    Feb 19, 2021 at 17:36

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service and acknowledge you have read our privacy policy.

Not the answer you're looking for? Browse other questions tagged or ask your own question.