How to assign class weights in the loss function in BCEloss

I have a 2d tensor of the shape (batch_size, 500), 500 is no of voice frames taken at a time and each of these frames has two labels i.e. 0 and 1, 0 denoting silence while 1 denotes speech.

tensor([[0., 0., 0.,  ..., 0., 0., 0.],
        [0., 0., 0.,  ..., 0., 0., 0.],
        [0., 0., 0.,  ..., 0., 0., 0.],
        ...,
        [0., 0., 0.,  ..., 1., 0., 0.],
        [0., 0., 0.,  ..., 0., 0., 0.],
        [0., 0., 1.,  ..., 0., 0., 0.]], device='cuda:0')

After passing the input features into my BiLSTM model from this is the output tensor

tensor([[0.5058, 0.5088, 0.5075,  ..., 0.5071, 0.5079, 0.5057],
        [0.5010, 0.4988, 0.4984,  ..., 0.5046, 0.5041, 0.5022],
        [0.5081, 0.5079, 0.5069,  ..., 0.4985, 0.4982, 0.4992],
        ...,
        [0.5064, 0.5104, 0.5117,  ..., 0.5039, 0.5040, 0.5041],
        [0.5049, 0.5075, 0.5079,  ..., 0.5178, 0.5174, 0.5162],
        [0.4936, 0.4948, 0.4970,  ..., 0.5033, 0.5038, 0.5041]],
       device='cuda:0', grad_fn=<SqueezeBackward0>)

Now, from the above tensor i am assigning 1 if the value is above 0.5 and 0 otherwise, so that each of the 500 frames has 1 or 0 assigned to it. After that i am calculating the BCEloss and then back propagating. But the issue is the input labels are unbalanced, so i want to assign class_weight to the labels during calculating the BCEloss.

After obtaining the class_weight from compute_class_weight module in sklearn, i am getting class_weights as array([0.59432247, 3.15048184]). But when i pass this as a tensor to the BCEloss i am getting an error.
RuntimeError: The size of tensor a (500) must match the size of tensor b (2) at non-singleton dimension 1
This is probably happening because the output tensor size is (batch_size, 500) while i have class_weights for two labels [0, 1].

Can anyone help me as to how i can assign the class_weights during training. Any kind of help would be greatly appreciated.

Thank You

Hi, Have you solved the issue? I am facing the same problem after passing the weight.
Thank you