shorten def

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