How to read the predicted label of a Neural Netowork with Cross Entropy Loss? Pytorch

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):
        super().__init__()
        # 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)
            loss.backward()
            optimizer.step()
            optimizer.zero_grad()
        # Validation phase
        result = evaluate(model, val_loader)
        model.epoch_end(epoch, result, epochs)
        history.append(result)
    return history

loss_value = evaluate(model, valid_dl)

#model=WineQuality()
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)
print(prediction)
output = model(input.unsqueeze(0))
_,pred = output.max(1)
print(pred)

And the output is the following:

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

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.

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

for i in range(num_input):
    print(classes[predicted[i]])

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

# print class name
print(classes[pred[0]])  # <- should be added