I am trying to assign different weights to different classes, so I have modified my loss criterion as such:
I had to convert the weight tensor to double torch.DoubleTensor(weight) since my model is already moved to double().
Am I doing this correctly ?
weights = [0.4,0.8,1.0]

class_weights = torch.DoubleTensor(weights).cuda()

criterion = nn.CrossEntropyLoss(weight=class_weights)