Correct weights assignment for CrossEntropyLoss and cuda()

Hi guys, below is simple example of neuralnetwork on Pytorch.
My dataset is very unbalanced (90% of class 0 and 10% of class 1).
As I learned on this forum, the best way to deal with is is to use “weight” parameter in CrossEntropyLoss.
I have to questoons:

  1. Should I input weights as [0.1, 0.9] or [0.9, 0.1]. How to check that weight is assigned to correct label?
  2. Do we need to use .cuda() for weights? If I use it, programme gives me an error:
    class_weights = torch.FloatTensor(weights).cuda()

Error:
Expected object of type torch.FloatTensor but found type torch.cuda.FloatTensor for argument #3 ‘weight’

Thank you for your help!

model = nn.Sequential(nn.Linear(16, 12),
                  nn.ReLU(),
                  nn.Linear(12, 4),
                  nn.ReLU(),
                  nn.Linear(4, 2))

weights = [0.1, 0.9]
class_weights = torch.FloatTensor(weights)
criterion = nn.CrossEntropyLoss(weight=class_weights)
optimizer = optim.RMSprop(model.parameters(), lr=0.003, momentum=0.9)

epochs = 5
for e in range(epochs):
    running_loss = 0
    X = Variable(torch.Tensor(trn_x).float())
    Y = Variable(torch.Tensor(trn_y).long())
    optimizer.zero_grad()
    
    output = model.forward(X)
    loss = criterion(output, Y)
    loss.backward()
    optimizer.step()
    
    running_loss += loss.item()
    print(f"Training loss: {running_loss}")

You should increase the weight of the minority class.
Based on the class distribution, [0.1, 0.9] should work better than the other way around.
If you would like to check for correctness, you could create dummy examples misclassifying one specific class and see, how the loss is scaled up or down.

If you don’t use the GPU for your model, you don’t have to push the weights onto the device either.