From 02982733ec3f9ca5da2011d77a3721787c1daf14 Mon Sep 17 00:00:00 2001 From: Wing Lian Date: Tue, 14 May 2024 08:17:30 -0400 Subject: [PATCH] fix attention mask collation (#1603) --- src/axolotl/utils/collators.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/src/axolotl/utils/collators.py b/src/axolotl/utils/collators.py index f0a1fb126..26c7fa9f3 100644 --- a/src/axolotl/utils/collators.py +++ b/src/axolotl/utils/collators.py @@ -229,9 +229,8 @@ class PretrainingBatchSamplerDataCollatorForSeq2Seq(DataCollatorForSeq2Seq): if feature == "attention_mask": if self.multipack_attn: arrays = [ - (i + 1) * np.array(item[feature]) + (i + 1) * np.array(item) for i, item in enumerate(features[feature]) - if feature in item ] else: arrays = [(1) * np.array(item) for item in features[feature]]