Fix IndexError for dec-only models in attribute-context

inseq-team · Mar 13, 2024 · ac3d18e · ac3d18e
1 parent 66f3f33
commit ac3d18e
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/inseq/models/decoder_only.py b/inseq/models/decoder_only.py
@@ -119,7 +119,7 @@ def enrich_step_output(
             contrast_target_ids = contrast_batch.target_ids[:, contrast_aligned_idx]
             step_output.target = join_token_ids(
                 tokens=target_tokens,
-                ids=attribution_model.convert_ids_to_tokens(contrast_target_ids),
+                ids=attribution_model.convert_ids_to_tokens(contrast_target_ids, skip_special_tokens=False),
                 contrast_tokens=attribution_model.convert_ids_to_tokens(
                     contrast_target_ids[None, ...], skip_special_tokens=False
                 ),