Sorry, I was talking about not being able to see the actual implementation of torch.binary_cross_entropy_with_logits. Unfortunately, passing pos_weight as a list of weights for each class doesn’t work. Here is an example:
# Random input and target of shape (N, C, H, W)
# Inputs are logits and target is binary
input = torch.rand(1, 3, 2, 2)
target = torch.round(torch.rand(1, 3, 2, 2))

criterion = BCEWithLogitsLoss(pos_weight=torch.tensor([1., 1., 1.]))
loss = criterion(input, target)

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

I also tried the channel_last variant with tensor shaped of (N, H, W, C). That works but outputs a different result than expected. It is hard to understand the behavior of pos_weight in this segmentation setting.