Save the unnormalised `lambda`, then normalise it afterward for safe softmaxing.
Save the unnormalised
lambda, then normalise it afterward for safe softmaxing.