fixing s2_loss annotate bug

FlamTeam · Jan 11, 2024 · f2e0496 · f2e0496
1 parent 79441d1
commit f2e0496
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/flamedisx/lxe_blocks/s2_loss.py b/flamedisx/lxe_blocks/s2_loss.py
@@ -45,7 +45,7 @@ def _annotate(self, d):
 
         mle = d['s2_raw' + '_mle'] = \
             (d['s2_raw_after_loss' + '_mle'] / s2_survival_probability).clip(0, None)
-        s = d['s2_raw'] * s2_survival_probability*(1-s2_survival_probability)
+        s = d['s2_raw' + '_mle'] * s2_survival_probability*(1-s2_survival_probability)
         scale = mle**0.5 * s / s2_survival_probability
 
         for bound, sign, intify in (('min', -1, np.floor),