Solved by ptrblck in post #2 
              
              
                nn.Softmax is an nn.Module so you would either have to create the module first and call it with output or use F.softmax instead: 
pred = torch.nn.Softmax(dim=1)(output)
# or
pred = F.softmax(output, dim=1)