Fix: prevent variance underflow causing NaN losses

[Kapadiaaa]

Problem
Softplus activation can produce zero values for very negative inputs, leading to division-by-zero and NaN losses during training.

Solution
Added a clamp (min=1e-6) to the predicted standard deviation to ensure numerical stability.

Impact

• Prevents NaN crashes
• Ensures stable probabilistic training

[kshirajahere]

Hey isnt there a draft PR #527 which resolves this. Given that overlap, I’d suggest closing this in favor of #527, unless there’s some unique piece here that should be folded into that branch.

[Kapadiaaa]

Thanks for pointing that out. I will review PR #527 and see if there’s anything additional I can contribute or improve there.