Hi, I am doing segmantic segmentation with large class imbalances(5 classes). So I am passing in a weight array into my loss  as in loss_fn = torch.nn.CrossEntropyLoss(weight=loss_weights)
Now since, it is very hard to assign these weights, I am trainning for 200 epochs and then setting the loss_weights then as trainable parameters. But when I do this I get the error :
RuntimeError: the derivative for ‘weight’ is not implemented
How can I get around this and any other suggestions to deal with such class imbalance?