Hello,
reduction=mean in BCELoss means average in pixel dimension, the result you get is a scalar.
If you want to sum loss in pixel dimention and perform averaging in batch dim, you could set reduction='none' and manually add sum and average operation.
criterion = nn.BCEWithLogitsLoss(reduction='none')
...
loss = criterion(logits, target).sum(dim=(1,2,3)).mean(dim=0)
loss.backward()