Passing the weights to CrossEntropyLoss correctly

I think we don’t need to normalize the weights of classes, using invert class sampling count is okay.