Fix hhbtag score handling (@taaaeen).

uhh-cms · Jan 13, 2025 · f30eb56 · f30eb56
1 parent 2efea55
commit f30eb56
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/hbt/production/hhbtag.py b/hbt/production/hhbtag.py
@@ -99,11 +99,11 @@ def split(where):
     even_mask = ak.to_numpy((events[event_mask].event % 2) == 0)
     if ak.sum(even_mask):
         input_features_even = split(even_mask)
-        scores_even = self.hhbtag_model_even(input_features_even)[0].numpy()
+        scores_even = self.hhbtag_model_even(input_features_even).numpy()
         scores[even_mask] = scores_even
     if ak.sum(~even_mask):
         input_features_odd = split(~even_mask)
-        scores_odd = self.hhbtag_model_odd(input_features_odd)[0].numpy()
+        scores_odd = self.hhbtag_model_odd(input_features_odd).numpy()
         scores[~even_mask] = scores_odd
 
     # remove the scores of padded jets