Hello,
I am trying to use BCEWithLogitsLoss() with additional dimensions but it crashes. The error messages are different in cases the additional axis has a size of 1 or 2:
batch_size = 8
groups = 1
num_classes = 10
pos_weight = (num_classes-1.0)*torch.ones([num_classes])
criterion = nn.BCEWithLogitsLoss(pos_weight=pos_weight)
x = torch.rand(batch_size, num_classes, groups)
print(x.shape)
# print(x)
y = torch.empty(batch_size, dtype=torch.long).random_(num_classes)
print(y.shape)
y = y.unsqueeze(1).expand(-1,groups)
print(y.shape)
y = F.one_hot(y, num_classes).type_as(x).transpose(1, 2)
print(y.shape)
# print(y)
criterion(x,y)

torch.Size([8, 10, 1])
torch.Size([8])
torch.Size([8, 1])
torch.Size([8, 10, 1])

---------------------------------------------------------------------------
RuntimeError                              Traceback (most recent call last)
<ipython-input-36-5314aeccd9bc> in <module>
     14 print(y.shape)
     15 # print(y)
---> 16 criterion(x,y)

/srv/storage/irim@storage1.lille.grid5000.fr/anaconda3/lib/python3.7/site-packages/torch/nn/modules/module.py in __call__(self, *input, **kwargs)
    530             result = self._slow_forward(*input, **kwargs)
    531         else:
--> 532             result = self.forward(*input, **kwargs)
    533         for hook in self._forward_hooks.values():
    534             hook_result = hook(self, input, result)

/srv/storage/irim@storage1.lille.grid5000.fr/anaconda3/lib/python3.7/site-packages/torch/nn/modules/loss.py in forward(self, input, target)
    599                                                   self.weight,
    600                                                   pos_weight=self.pos_weight,
--> 601                                                   reduction=self.reduction)
    602 
    603 

/srv/storage/irim@storage1.lille.grid5000.fr/anaconda3/lib/python3.7/site-packages/torch/nn/functional.py in binary_cross_entropy_with_logits(input, target, weight, size_average, reduce, reduction, pos_weight)
   2124         raise ValueError("Target size ({}) must be the same as input size ({})".format(target.size(), input.size()))
   2125 
-> 2126     return torch.binary_cross_entropy_with_logits(input, target, weight, pos_weight, reduction_enum)
   2127 
   2128 

RuntimeError: output with shape [8, 10, 1] doesn't match the broadcast shape [8, 10, 10]


batch_size = 8
groups = 2
num_classes = 10
pos_weight = (num_classes-1.0)*torch.ones([num_classes])
criterion = nn.BCEWithLogitsLoss(pos_weight=pos_weight)
x = torch.rand(batch_size, num_classes, groups)
print(x.shape)
# print(x)
y = torch.empty(batch_size, dtype=torch.long).random_(num_classes)
print(y.shape)
y = y.unsqueeze(1).expand(-1,groups)
print(y.shape)
y = F.one_hot(y, num_classes).type_as(x).transpose(1, 2)
print(y.shape)
# print(y)
criterion(x,y)


torch.Size([8, 10, 2])
torch.Size([8])
torch.Size([8, 2])
torch.Size([8, 10, 2])

---------------------------------------------------------------------------
RuntimeError                              Traceback (most recent call last)
<ipython-input-37-f22703bfb833> in <module>
     14 print(y.shape)
     15 # print(y)
---> 16 criterion(x,y)

/srv/storage/irim@storage1.lille.grid5000.fr/anaconda3/lib/python3.7/site-packages/torch/nn/modules/module.py in __call__(self, *input, **kwargs)
    530             result = self._slow_forward(*input, **kwargs)
    531         else:
--> 532             result = self.forward(*input, **kwargs)
    533         for hook in self._forward_hooks.values():
    534             hook_result = hook(self, input, result)

/srv/storage/irim@storage1.lille.grid5000.fr/anaconda3/lib/python3.7/site-packages/torch/nn/modules/loss.py in forward(self, input, target)
    599                                                   self.weight,
    600                                                   pos_weight=self.pos_weight,
--> 601                                                   reduction=self.reduction)
    602 
    603 

/srv/storage/irim@storage1.lille.grid5000.fr/anaconda3/lib/python3.7/site-packages/torch/nn/functional.py in binary_cross_entropy_with_logits(input, target, weight, size_average, reduce, reduction, pos_weight)
   2124         raise ValueError("Target size ({}) must be the same as input size ({})".format(target.size(), input.size()))
   2125 
-> 2126     return torch.binary_cross_entropy_with_logits(input, target, weight, pos_weight, reduction_enum)
   2127 
   2128 

RuntimeError: The size of tensor a (10) must match the size of tensor b (2) at non-singleton dimension 2

In both cases, I am unable to understand the error message. The x and y tensors have the same shape and the documentation says that there can be any number of additional dimensions after the batch one. Can anyone tell me what I got wrong?
Thanks.
Best regards,
Georges Quénot.