Passing the weights to CrossEntropyLoss correctly