Fix: Noise-noise connections should not count as true edges (#446)

gnn-tracking · Oct 2, 2023 · 2924c13 · 2924c13
1 parent 7ed880c
commit 2924c13
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/src/gnn_tracking/models/graph_construction.py b/src/gnn_tracking/models/graph_construction.py
@@ -318,7 +318,7 @@ def forward(self, data: Data) -> Data:
         )
         y: T = (  # type: ignore
             data.particle_id[edge_index[0]] == data.particle_id[edge_index[1]]
-        )
+        ) & (data.particle_id[edge_index[0]] > 0)
         if not self._ml or not self.hparams.use_embedding_features:
             x = data.x
         else: