fixing s2_loss annotate bug

FlamTeam · Jan 11, 2024 · 79441d1 · 79441d1
1 parent b21805c
commit 79441d1
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
@@ -44,7 +44,7 @@ def _annotate(self, d):
         s2_survival_probability = self.gimme_numpy('s2_survival_p')
 
         mle = d['s2_raw' + '_mle'] = \
-            (d['s2_raw_after_loss_' + '_mle'] / s2_survival_probability).clip(0, None)
+            (d['s2_raw_after_loss' + '_mle'] / s2_survival_probability).clip(0, None)
         s = d['s2_raw'] * s2_survival_probability*(1-s2_survival_probability)
         scale = mle**0.5 * s / s2_survival_probability