I am dealing with imbalance labels in my image segmentation. The model output -> 1, 5, 224, 224 is used for the prediction of two tasks task_a = output[1, [0, 1, 2], :,:] and task_b = output[1, [0, 1], :,:] . I am using the CrossEntropy loss and when I pass weight parameter `CrossEntropyLoss(weight = [0.5, 0.3, 0.1, 0.6, 0.5]), the runtime error is outputted
  File "/home/pc-man/.local/lib/python3.8/site-packages/torch/nn/modules/module.py", line 727, in _call_impl
    result = self.forward(*input, **kwargs)
  File "/home/pc-man/Documents/GAN/Code/main/losses/dice.py", line 84, in forward
    ce_loss = self.cross_entroy(predictions, labels)
  File "/home/pc-man/.local/lib/python3.8/site-packages/torch/nn/modules/module.py", line 727, in _call_impl
    result = self.forward(*input, **kwargs)
  File "/home/pc-man/.local/lib/python3.8/site-packages/torch/nn/modules/loss.py", line 961, in forward
    return F.cross_entropy(input, target, weight=self.weight,
  File "/home/pc-man/.local/lib/python3.8/site-packages/torch/nn/functional.py", line 2468, in cross_entropy
    return nll_loss(log_softmax(input, 1), target, weight, None, ignore_index, None, reduction)
  File "/home/pc-man/.local/lib/python3.8/site-packages/torch/nn/functional.py", line 2266, in nll_loss
    ret = torch._C._nn.nll_loss2d(input, target, weight, _Reduction.get_enum(reduction), ignore_index)
RuntimeError: weight tensor should be defined either for all or no classes