Solved by ptrblck in post #2 
              
              
                The dim argument defines which dimension should be used to calculate the log softmax, i.e. in which dimension the class logits are located. 
Have a look at this small example using softmax: 
x = torch.randn(5, 3)
x0 = F.softmax(x, dim=0)
print(x0)
> tensor([[0.1313, 0.0170, 0.4122],
        [0.0167,…