From 155be7510778b56329f6723ab7abda9e4596fe76 Mon Sep 17 00:00:00 2001 From: BioMikeUkr Date: Mon, 2 Feb 2026 16:59:01 +0200 Subject: [PATCH] indices out of range fixed --- gliner/modeling/span_rep.py | 8 ++------ 1 file changed, 2 insertions(+), 6 deletions(-) diff --git a/gliner/modeling/span_rep.py b/gliner/modeling/span_rep.py index 954ac0f5..9f737169 100644 --- a/gliner/modeling/span_rep.py +++ b/gliner/modeling/span_rep.py @@ -374,14 +374,10 @@ def extract_elements(sequence, indices): Returns: torch.Tensor: Extracted elements of shape [B, K, D]. """ - D = sequence.size(-1) - - # Expand indices to [B, K, D] + B, L, D = sequence.size() + indices = torch.clamp(indices, 0, L - 1) expanded_indices = indices.unsqueeze(2).expand(-1, -1, D) - - # Gather the elements extracted_elements = torch.gather(sequence, 1, expanded_indices) - return extracted_elements