Losses from logits
A model's raw scores are called logits. Turning them into probabilities first and taking the log afterwards loses precision: sigmoid(40) rounds to exactly 1.0 in float32, so log(1 - p) is minus infinity.
import torch
import torch.nn.functional as F
z = torch.tensor([40.0, -40.0])
y = torch.tensor([0.0, 1.0])
p = torch.sigmoid(z)
naive = -(y * p.log()
+ (1 - y) * (1 - p).log())
print(naive) # inf where p rounded to 1
print(F.binary_cross_entropy_with_logits(
z, y, reduction="none"))binary_cross_entropy_with_logits and cross_entropy take logits and rearrange the formula so nothing rounds away. For two classes the stable form is max(z, 0) - z * y + log(1 + exp(-abs(z))).
When one class is rare, the model can score well by ignoring it. nn.CrossEntropyLoss(weight=w) multiplies each example's loss by its class's weight; a common choice is n / (classes * count), so every class carries the same total weight.
Your task
Write bce_logits(z, y): the mean binary cross-entropy of logits z against targets y of 0s and 1s, with the stable formula and without sigmoid. Write class_weights(labels, classes): a float tensor where class c gets len(labels) / (classes * count_c), and a class that never occurs gets 0.
The explanation opens here once every check passes.