Clamp before .sqrt() in newton_2d()
@@ -203,7 +203,7 @@ def newton_step_2d(loss, x, trust_radius=None):
# regularize to keep update within ball of given trust_radius
detH = H[..., 0, 0] * H[..., 1, 1] - H[..., 0, 1] * H[..., 1, 0]
mean_eig = (H[..., 0, 0] + H[..., 1, 1]) / 2
- min_eig = mean_eig - (mean_eig ** 2 - detH).sqrt()
+ min_eig = mean_eig - (mean_eig ** 2 - detH).clamp(min=0).sqrt()
regularizer = (g.pow(2).sum(-1).sqrt() / trust_radius - min_eig).clamp_(min=1e-8)
warn_if_nan(regularizer, 'regularizer')
H = H + regularizer.unsqueeze(-1).unsqueeze(-1) * H.new_tensor(torch.eye(2))
