Hotfix from_step_attribution for decoder self-attention

inseq-team · Jan 17, 2024 · 09c8481 · 09c8481
1 parent f434192
commit 09c8481
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/inseq/data/attribution.py b/inseq/data/attribution.py
@@ -274,7 +274,7 @@ def from_step_attributions(
                     out_seq_scores = [attr.sequence_scores[seq_score_name][i, ...] for i in range(num_sequences)]
                 else:
                     out_seq_scores = get_sequences_from_batched_steps(
-                        [att.sequence_scores[seq_score_name] for att in attributions], padding_dims=[2], stack_dim=3
+                        [att.sequence_scores[seq_score_name] for att in attributions], padding_dims=[1]
                     )
                 for seq_id in range(num_sequences):
                     seq_scores[seq_id][seq_score_name] = remove_pad_fn(out_seq_scores, sources, targets, seq_id)