add missing abs to track label

RasmusOrsoe · May 20, 2024 · b31aa42 · b31aa42
1 parent 2efd732
commit b31aa42
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/src/graphnet/training/labels.py b/src/graphnet/training/labels.py
@@ -102,5 +102,5 @@ def __init__(
 
     def __call__(self, graph: Data) -> torch.tensor:
         """Compute label for `graph`."""
-        label = (graph[self._pid_key] == 14) & (graph[self._int_key] == 1)
+        label = (torch.abs(graph[self._pid_key]) == 14) & (graph[self._int_key] == 1)
         return label.type(torch.int)