I am trying to train the weight used in CrossEntropyLoss, and my code is something like this:
...
loss_weights = nn.Parameter(nn.Softmax()(torch.randn(n_classes, device=device, requires_grad=True)))
criterion = nn.CrossEntropyLoss(weight=loss_weights, size_average=False)
optimizer_ft = optim.SGD(list(mdl.parameters()) + list(loss_weights), lr=0.001, momentum=0.9)
...

But is report ‘can’t optimized non-leaf Tensor’ error.
How is that possible to do this?