Merge pull request #725 from RasmusOrsoe/fix_track_label

add missing `abs` to `Track` label
graphnet-team · May 28, 2024 · e5b9450 · e5b9450
2 parents fc5d955 + 1ee2b53
commit e5b9450
Showing 1 changed file with 3 additions and 2 deletions.
diff --git a/src/graphnet/training/labels.py b/src/graphnet/training/labels.py
@@ -102,5 +102,6 @@ def __init__(
 
     def __call__(self, graph: Data) -> torch.tensor:
         """Compute label for `graph`."""
-        label = (graph[self._pid_key] == 14) & (graph[self._int_key] == 1)
-        return label.type(torch.int)
+        is_numu = torch.abs(graph[self._pid_key]) == 14
+        is_cc = graph[self._int_key] == 1
+        return (is_numu & is_cc).type(torch.int)