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)