Solved by KFrank in post #3 
              
              
                Hi Berkay! 

Perform your s, (1 - s) weighting in “log space” so that you work 
directly with the log-probabilities and only have to call log_softmax(), 
with its better numerical stability. 
That is, because: 
log (s * prob) = log (s) + log_prob,

just add log (s) (and log (1 - s)) to your results …