Where did you modify the code? In your first post or the github repo?
I’m not sure about this line:
label = label[:, :, np.newaxis]

Are you trying to add a batch dimension to your target?
If so, the batch dimension should be at dim0, i.e. the first dimension.
Make sure your label is of type long.
Here is a small code example using F.cross_entropy:
batch_size = 5
nb_classes = 3

output = torch.randn(batch_size, nb_classes)
target = torch.empty(batch_size, dtype=torch.long).random_(nb_classes)

loss = F.cross_entropy(output, target)