Merge pull request #19 from rkansal47/main

Minor FPND and utils.gen_jet_corrections bug fixes
jet-net · Sep 4, 2022 · 11cffb6 · 11cffb6
2 parents 82129fd + acdc6f4
commit 11cffb6
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 2 deletions.
diff --git a/jetnet/evaluation/gen_metrics.py b/jetnet/evaluation/gen_metrics.py
@@ -273,7 +273,7 @@ def fpnd(
     sigma1 = _eval_module.fpnd_dict[dataset_name][num_particles][jet_type]["sigma"]
 
     # run inference and store activations
-    jets_loaded = DataLoader(jets[:, :, : _eval_module.fpnd_dict["NUM_SAMPLES"]], batch_size)
+    jets_loaded = DataLoader(jets[: _eval_module.fpnd_dict["NUM_SAMPLES"]], batch_size)
 
     logging.info(f"Calculating ParticleNet activations with batch size: {batch_size}")
     activations = []

diff --git a/jetnet/utils/utils.py b/jetnet/utils/utils.py
@@ -188,4 +188,4 @@ def gen_jet_corrections(
     if zero_neg_pt:
         jets[:, :, pt_index][jets[:, :, pt_index] < 0] = 0
 
-    return jets[:, :, :-1], mask if ret_mask_separate else jets
+    return (jets[:, :, :-1], mask) if ret_mask_separate else jets