From 82d04ea060fd12fa15ea4ad9f6d1ceacb50cdc03 Mon Sep 17 00:00:00 2001
From: bursteratom <sunny19981005@outlook.com>
Date: Thu, 13 Feb 2025 00:11:45 -0500
Subject: [PATCH] test v2batch w/ flex attn

---
 src/axolotl/core/trainer_builder.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/axolotl/core/trainer_builder.py b/src/axolotl/core/trainer_builder.py
index 784b3b697..1ca992b83 100755
--- a/src/axolotl/core/trainer_builder.py
+++ b/src/axolotl/core/trainer_builder.py
@@ -830,7 +830,7 @@ class HFCausalTrainerBuilder(TrainerBuilderBase):
                 kwargs.pop("max_length")
         elif use_batch_sampler_collator:
             if self.cfg.flex_attention is True:
-                collator = FlexBatchSamplerDataCollatorForSeq2Seq
+                collator = V2BatchSamplerDataCollatorForSeq2Seq
             elif self.cfg.model_config_type in SUPPORTED_MULTIPACK_MODEL_TYPES:
                 collator = V2BatchSamplerDataCollatorForSeq2Seq
             elif (