From d23aab2578ef0df901ae9584240fa68b4d4ea735 Mon Sep 17 00:00:00 2001 From: Elihei2 Date: Mon, 1 Jun 2026 14:33:07 +0200 Subject: [PATCH] fix(model): guard prediction against empty tiles Two robustness fixes that surface on sparse/empty tiles (common at scale, e.g. MERSCOPE): - LitISTEncoder.predict_step: scatter_max returns -1 for transcripts with no candidate boundary; main only checked `max_idx < dst.shape[0]` so a -1 was treated as valid and indexed wrongly. Add the `max_idx >= 0` guard. - Positional2dEmbedder.forward: return zeros for empty `pos`/`batch` instead of crashing on min/max over an empty tensor; cast batch to long. What to review: the one-line `valid` guard in lightning_model.py and the two empty-tensor early-returns in ist_encoder.py. No behavior change on non-empty input. --- src/segger/models/ist_encoder.py | 6 ++++++ src/segger/models/lightning_model.py | 2 +- 2 files changed, 7 insertions(+), 1 deletion(-) diff --git a/src/segger/models/ist_encoder.py b/src/segger/models/ist_encoder.py index 04d4ef9..2b73c3e 100644 --- a/src/segger/models/ist_encoder.py +++ b/src/segger/models/ist_encoder.py @@ -59,10 +59,16 @@ def forward( pos: torch.Tensor, batch: Optional[torch.Tensor] = None, ) -> torch.Tensor: + if pos.numel() == 0: + return pos.new_zeros((pos.shape[0], self.dim * 2)) + if batch is None: pos = pos - pos.min(dim=0).values pos = pos / pos.max(dim=0).values else: + if batch.numel() == 0: + return pos.new_zeros((pos.shape[0], self.dim * 2)) + batch = batch.to(torch.long) # normalize per batch mins = torch.zeros((batch.max()+1, 2), device=pos.device) maxs = torch.zeros((batch.max()+1, 2), device=pos.device) diff --git a/src/segger/models/lightning_model.py b/src/segger/models/lightning_model.py index 3578a16..48c3657 100644 --- a/src/segger/models/lightning_model.py +++ b/src/segger/models/lightning_model.py @@ -283,7 +283,7 @@ def predict_step( dim_size=batch['tx'].num_nodes, ) # Filter by similarity - valid = max_idx < dst.shape[0] + valid = (max_idx >= 0) & (max_idx < dst.shape[0]) if min_similarity is not None: valid &= max_sim >= min_similarity