deJQK:

loss_weights = nn.Parameter(nn.Softmax()(torch.randn(n_classes, device=device, requires_grad=True)))


The problem with this statement is that a leaf tensor is being created (torch.randn(..., requires_grad=True)) and then it is being hidden because nn.Softmax() returns a new tensor.
To make this work, try something like:
initial_weights = nn.Softmax()(torch.randn(n_classes, device=device)
loss_weights = torch.zeros_like(initial_weights, requires_grad=True)
# Don't record the following operation in autograd
with torch.no_grad():
    loss_weights.copy_(initial_weights)

and then proceed to optimize loss_weights