How to check corresponding labels for predictions/softmax values?

1

Suppose that I have a tensor produced by F.softmax(nn.Linear()) of shape [64, 2], where 64 is the batch size and 2 is the number of classes;

tensor([[0.5382, 0.4618],
        [0.5180, 0.4820],
        [0.5294, 0.4706],
        [0.5272, 0.4728],
        [0.5307, 0.4693],
        [0.5291, 0.4709],
        [0.5215, 0.4785],
        [0.5300, 0.4700],
        [0.5330, 0.4670],

And that my label/class can either be 0 or 1.
How can I know which probability corresponds to which label? I suppose that at training time, I need to pick the probability value of the true tag and feed that to the loss function.

2

For the first sample, you have 0.5382 probability for class 0 and 0.4618 probability for class 1.