I was following this tutorial 22 and was confused by the way we make final prediction. Let’s say the model is defined as
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!