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.