Hello, I have read several topics about setting the weight to a loss, but I have some interesting to me question.

So I have a binary segmentation problem, with classes 0 – background, and 1 – buildings. And they are unbalanced. I decided to set a weight for BCEWithLogits loss with `torch.tensor([0.3, 0.7])`

for class 0 and 1, respectively. But when I try to calculate the loss, script throws a Runtime error:

```
import torch
output = torch.randn(1, 2, 256, 256, requires_grad=True)
target = torch.randn(1, 2, 256, 256, requires_grad=False)
criterion = torch.nn.BCEWithLogitsLoss(pos_weight=torch.tensor([0.3, 0.7]))
# same when I use 'weight' insted pos_weight
loss = criterion(output, target)
```

**Error**

```
---------------------------------------------------------------------------
RuntimeError Traceback (most recent call last)
<ipython-input-9-0ea142ac673d> in <module>
3 target = torch.randn(1, 2, 256, 256, requires_grad=False)
4 criterion = torch.nn.BCEWithLogitsLoss(weight=torch.tensor([0.3, 0.7]))
----> 5 loss = criterion(output, target)
/opt/conda/lib/python3.7/site-packages/torch/nn/modules/module.py in __call__(self, *input, **kwargs)
548 result = self._slow_forward(*input, **kwargs)
549 else:
--> 550 result = self.forward(*input, **kwargs)
551 for hook in self._forward_hooks.values():
552 hook_result = hook(self, input, result)
/opt/conda/lib/python3.7/site-packages/torch/nn/modules/loss.py in forward(self, input, target)
615 self.weight,
616 pos_weight=self.pos_weight,
--> 617 reduction=self.reduction)
618
619
/opt/conda/lib/python3.7/site-packages/torch/nn/functional.py in binary_cross_entropy_with_logits(input, target, weight, size_average, reduce, reduction, pos_weight)
2433 raise ValueError("Target size ({}) must be the same as input size ({})".format(target.size(), input.size()))
2434
-> 2435 return torch.binary_cross_entropy_with_logits(input, target, weight, pos_weight, reduction_enum)
2436
2437
RuntimeError: The size of tensor a (256) must match the size of tensor b (2) at non-singleton dimension 3
```

But when I try only one tensor with shape 1, it turns out that it works

```
import torch
output = torch.randn(1, 2, 256, 256, requires_grad=True)
target = torch.randn(1, 2, 256, 256, requires_grad=False)
criterion = torch.nn.BCEWithLogitsLoss(pos_weight=torch.tensor([0.3]))
loss = criterion(output, target)
print(loss)
>>> tensor(0.2424, grad_fn=<BinaryCrossEntropyWithLogitsBackward>)
```

My question is why the first method didn’t work, and the second worked. And if the second method is the correct one, then which class will get correct weight to get better loss?

Thank you!