# Making prediction with argmax

**URL:** <https://discuss.pytorch.org/t/making-prediction-with-argmax/49526>\
**Category:** Uncategorized\
**Created:** [July 2, 2019, 8:45pm UTC](https://discuss.pytorch.org/t/making-prediction-with-argmax/49526 "2019-07-02T20:45:33Z")\
**Posts on this page:** 2\
**Page:** 1

<div class="post-metadata">

**Author:** ![hiepnguyen034](https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/hiepnguyen034/32/12872_2.png) [@hiepnguyen034](https://discuss.pytorch.org/u/hiepnguyen034)\
**Post date:** [July 2, 2019, 8:45pm UTC](https://discuss.pytorch.org/t/making-prediction-with-argmax/49526/1 "2019-07-02T20:45:33Z")

</div>

I was following [this tutorial](https://pytorch.org/tutorials/beginner/nn_tutorial.html) and was confused by the way we make final prediction. Let’s say the model is defined as

```auto
loss_func = F.cross_entropy

def get_model():
    model = Mnist_Logistic()
    return model, optim.SGD(model.parameters(), lr=lr)

model, opt = get_model()
print(loss_func(model(xb), yb))

for epoch in range(epochs):
    for i in range((n - 1) // bs + 1):
        start_i = i * bs
        end_i = start_i + bs
        xb = x_train[start_i:end_i]
        yb = y_train[start_i:end_i]
        pred = model(xb)
        loss = loss_func(pred, yb)

        loss.backward()
        opt.step()
        opt.zero_grad()

print(loss_func(model(xb), yb))

```

. Then, the final prediction can be achieved by calling `torch.argmax(model(test_data), dim=1`). This means that `y_pred=model(test_data)` will return a vector that has a probability value for each digit. Then, my question is, how can PyTorch know index `0` of y\_pred will correspond to the probability of the digit being `0`, index `1` will correspond to the probability of the digit being `1`, and so on? Since our `yb` is a one-dimensional vector that looks like `[5, 0, 4, ..., 8, 4, 8]` instead of a one-hot matrix, how can PyTorch know to link the index of `y_pred` with the correct label?

Thank you!

---

<div class="post-metadata">

**Author:** ![KFrank](https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/k/ecb155/32.png) [@KFrank](https://discuss.pytorch.org/u/KFrank)\
**Post date:** [July 2, 2019, 9:41pm UTC](https://discuss.pytorch.org/t/making-prediction-with-argmax/49526/2 "2019-07-02T21:41:32Z")

</div>

Hi Hiep!

> [@hiepnguyen034](#):
>
> …
> 
> ```auto
> loss_func = F.cross_entropy
> 
> ```
> 
> …
> 
> Then, the final prediction can be achieved by calling `torch.argmax(model(test_data), dim=1`). This means that `y_pred=model(test_data)` will return a vector that has a probability value for each digit. Then, my question is, how can PyTorch know index `0` of y\_pred will correspond to the probability of the digit being `0`, index `1` will correspond to the probability of the digit being `1`, and so on? Since our `yb` is a one-dimensional vector that looks like `[5, 0, 4, ..., 8, 4, 8]` instead of a one-hot matrix, how can PyTorch know to link the index of `y_pred` with the correct label?
> 
> Thank you!

The short answer is that `nn.functional.cross_entropy`  
one-hots your class labels for you.

The number-one rule is that the output of your network means  
whatever you train it to mean.

More directly to your question:

You are using `nn.functional.cross_entropy` as your loss  
function. Your `yb` are integer class labels, one per sample.

Conceptually, implicitly under the hood, `cross_entropy`  
one-hots your `yb`, implicitly softmaxes your `pred`, and then  
implicitly calculates the cross-entropy of “softmax (pred)” and  
“one-hot (yb)”.

The output of your model (the `y_pred`) should be understood  
as logits. They (implicitly) get turned into probabilities when  
`cross_entropy` (implicitly) softmaxes them. So you are training  
your model to output (for each sample) a vector of length nClass,  
where the value for index i is the logit (sort of like the probability)  
of that sample being of class i. (Finally, you take the `argmax` of  
your prediction vector. This finds the index of the logit with the  
largest value – that is the index that your model predicts as having  
the highest probability of being the class label, and you take this  
as being the predicted class label.)

“PyTorch knows to link the index of `y_pred` with the correct label”  
because you trained your network to do so.

Good Luck.

K. Frank