Hello!
I have tried using weighted cross entropy and my training script looks like this:
weights = [0.2939, 0.0105, 0.0005, 0.0015, 0.0002, 0.0015, 0.0551, 0.0181, 0.0453, 0.0135, 0.0037, 0.0049, 0.0024, 0.0370, 0.0007, 0.0193, 0.0022, 0.0100, 0.0100, 0.0066, 0.0007, 0.0152, 0.0223, 0.0027, 0.0010, 0.0000, 0.0017, 0.0010, 0.0169, 0.0015, 0.0419, 0.0012, 0.0007, 0.0027, 0.0167, 0.0056, 0.0000, 0.0759, 0.0010, 0.0049, 0.0054, 0.0120, 0.0007, 0.0086, 0.0022, 0.0010, 0.0091, 0.0032, 0.0005, 0.0010, 0.0071, 0.0174, 0.0069, 0.0029, 0.0015, 0.0015, 0.0269, 0.0010, 0.0191, 0.0659, 0.0556]
class_weights = torch.FloatTensor(weights)

import torch.optim as optim
optimizer = optim.Adam(model.parameters(), lr = 3e-5)
loss_CE = nn.CrossEntropyLoss(weight=class_weights)

BATCH_SIZE = 4
EPOCHS = 2


def train(model):
  model.train()
  for epoch in range(EPOCHS):
      for i in tqdm(range(0, len(img_train), BATCH_SIZE)): 
          batch_img_train = img_train[i:i+BATCH_SIZE].view(-1, 3, 112, 112)
          batch_mask_train = mask_train[i:i+BATCH_SIZE].view(-1, 1, 112, 112)
        
          model.zero_grad()

          outputs = model(batch_img_train)
          loss_ce= loss_CE(outputs, batch_mask_train.squeeze(1).long())
          loss = loss_ce
        
          loss.backward()
          optimizer.step()    # Does the update

      print(f"Epoch: {epoch}. Loss: {loss}")

train(model)


Am I doing it right?