Merged
Keep focal_loss gradients finite for fractional gamma #5563
- Bug
- With 0 < γ < 1, (1 − p)^γ has an infinite derivative at p = 1. As soon as a softmax probability saturated, the chain rule multiplied that infinity by zero and every logit of the pixel got a NaN gradient, including when the saturated class was the wrong one.
- Fix
- The double-where pattern: saturated entries take their value from a detached copy and their gradient from a safe base, so forward values are unchanged and the gradient becomes its true limit, 0. Clamp was avoided because its derivative at the bound differs across torch versions.
- Test
- Saturated rows on the correct and the wrong class, γ in {0, 0.25, 0.5, 0.75, 1, 2}, every reduction: loss and gradient must match the γ = 0 case. Plus a torch.compile vs eager check.